train.py 文件源码

def save(images, depths, predict_depths, global_step, target_path, batch_number=None, mode='train'):

    output_dir = os.path.join(target_path, str(global_step))

    if not gfile.Exists(output_dir):
        gfile.MakeDirs(output_dir)
    for i, (image, depth, predict_depth) in enumerate(zip(images, depths, predict_depths)):
        if(batch_number == None):
            image_name = "%s/%05d_rgb.png" % (output_dir, i)
            depth_name = "%s/%05d_depth.png" % (output_dir, i)
            predict_depth_name = "%s/%05d_predict.png" % (output_dir, i)
        else:
            image_name = "%s/%d_%05d_rgb.png" % (output_dir, batch_number, i)
            depth_name = "%s/%d_%05d_depth.png" % (output_dir, batch_number, i)
            predict_depth_name = "%s/%d_%05d_predict.png" % (output_dir, batch_number, i)


        pilimg = Image.fromarray(np.uint8(image))
        pilimg.save(image_name)

        depth = depth.transpose(2, 0, 1)
        if np.max(depth) != 0:
            ra_depth = (depth/np.max(depth))*255.0
        else:
            ra_depth = depth*255.0
        depth_pil = Image.fromarray(np.uint8(ra_depth[0]), mode="L")
        depth_pil.save(depth_name)

        predict_depth = predict_depth.transpose(2, 0, 1)
        if np.max(predict_depth) != 0:
            ra_depth = (predict_depth/np.max(predict_depth))*255.0
        else:
            ra_depth = predict_depth*255.0
        depth_pil = Image.fromarray(np.uint8(ra_depth[0]), mode="L")
        depth_pil.save(predict_depth_name)