def main(args, config): cuda = cnn_utils.check_cuda(config) model = cnn_utils.load_model_and_weights(args, config) if cuda: model = model.cuda() file_path = os.path.join(config['PATH']['hdf5_dir'], config['PATH']['hdf5_name']) with h5py.File(file_path, mode='r', libver='latest') as hdf5_file: overall_psnr_accum = (0, 0, 0) overall_ssim_accum = (0, 0, 0) for sample_num in range(args.nSamples): p1, s1 = do_one_demo(args, config, hdf5_file, model, sample_num, cuda) overall_psnr_accum = welford.update(overall_psnr_accum, p1) overall_ssim_accum = welford.update(overall_ssim_accum, s1) if args.nSamples > 1: psnr_mean, psnr_var, _ = welford.finalize(overall_psnr_accum) ssim_mean, ssim_var, _ = welford.finalize(overall_ssim_accum) print("\nOverall cnn psnr average {:5f}, stddev {:5f}".format( psnr_mean, math.sqrt(psnr_var))) print("Overall cnn ssim average {:5f}, stddev {:5f}".format( ssim_mean, math.sqrt(ssim_var))) #Ground truth possible """
def main(pixel_dim, clip, num_random, plane): #Setup app = inviwopy.app network = app.network cam = network.EntryExitPoints.camera cam.nearPlane = 6.0 cam.farPlane = 1000.0 canvases = inviwopy.app.network.canvases for canvas in canvases: canvas.inputSize.dimensions.value = ivec2(pixel_dim, pixel_dim) inviwo_utils.update() random_lfs = create_random_lf_cameras( num_random, (180, 35), 1, interspatial_distance=0.5, look_up = vec3(0, 1, 0)) time_accumulator = (0.0, 0.0, 0.0) for lf in random_lfs: if clip: _, clip_type = random_clip_lf(network, lf) elif plane: random_plane_clip(network, lf) time_taken = lf.view_array(cam, save=False, should_time=True) time_accumulator = welford.update( time_accumulator, time_taken) if clip: restore_clip(network, clip_type) mean, variance, _ = welford.finalize(time_accumulator) print("Time taken per grid, average {:4f}, std_dev {:4f}".format( mean, math.sqrt(variance)))
def main(num_samples): random.seed(time.time()) time_accumulator = (0, 0, 0) for _ in range(num_samples): last_time = time_one() time_accumulator = welford.update(time_accumulator, last_time) if num_samples > 1: mean_time, std_dev_time, _ = welford.finalize(time_accumulator) print("Overall time mean: {:4f}, stdev: {:4f}".format(mean_time, std_dev_time))
def main(args): with h5py.File(args.loc, 'r') as f: overall_psnr_accum = (0, 0, 0) overall_ssim_accum = (0, 0, 0) i = args.n psnr_accumulator = (0, 0, 0) ssim_accumulator = (0, 0, 0) group1 = f[args.group]["images"][i, :, :args.channels, ...] group2 = f[args.group]["warped"][i, :, :args.channels, ...] for j in range(64): im1_out_location = os.path.join(args.out_loc, "{}/gt{}.png".format(i, j)) im2_out_location = os.path.join(args.out_loc, "{}/warp{}.png".format(i, j)) im1 = group1[j] im2 = group2[j] im1 = np.swapaxes(im1, 0, 2) im1 = np.swapaxes(im1, 0, 1) im2 = np.swapaxes(im2, 0, 2) im2 = np.swapaxes(im2, 0, 1) common.save_numpy_image(array=im1, location=im1_out_location) common.save_numpy_image(array=im2, location=im2_out_location) psnr = evaluate.psnr(im1, im2) ssim = evaluate.ssim(im1, im2) psnr_accumulator = welford.update(psnr_accumulator, psnr) ssim_accumulator = welford.update(ssim_accumulator, ssim) if args.verbose: print("{};{};PSNR {:4f};SSIM {:4f}".format(i, j, psnr, ssim)) psnr_mean, psnr_var, _ = welford.finalize(psnr_accumulator) ssim_mean, ssim_var, _ = welford.finalize(ssim_accumulator) if args.verbose: print() print("{};psnr average {:5f};stddev {:5f}".format( i, psnr_mean, math.sqrt(psnr_var)) + ";ssim average {:5f};stddev {:5f}".format( ssim_mean, math.sqrt(ssim_var))) overall_psnr_accum = welford.update(overall_psnr_accum, psnr_mean) overall_ssim_accum = welford.update(overall_ssim_accum, ssim_mean)
def main(args, config, sample_index): cuda = cnn_utils.check_cuda(config) model = cnn_utils.load_model_and_weights(args, config) if cuda: model = model.cuda() model.eval() # Create output directory base_dir = os.path.join(config['PATH']['output_dir'], 'warped') if not os.path.isdir(base_dir): pathlib.Path(base_dir).mkdir(parents=True, exist_ok=True) save_dir = get_sub_dir_for_saving(base_dir) start_time = time.time() file_path = os.path.join(config['PATH']['hdf5_dir'], config['PATH']['hdf5_name']) with h5py.File(file_path, mode='r', libver='latest') as hdf5_file: depth_grp = hdf5_file['val']['disparity'] SNUM = sample_index depth_images = torch.squeeze( torch.tensor(depth_grp['images'][SNUM], dtype=torch.float32)) colour_grp = hdf5_file['val']['colour'] colour_images = torch.tensor(colour_grp['images'][SNUM], dtype=torch.float32) sample = { 'depth': depth_images, 'colour': colour_images, 'grid_size': depth_images.shape[0] } warped = data_transform.center_normalise(sample) im_input = warped['inputs'].unsqueeze_(0) if cuda: im_input = im_input.cuda() output = model(im_input) time_taken = time.time() - start_time print("Time taken was {:4f}s".format(time_taken)) grid_size = 64 psnr_accumulator = (0, 0, 0) ssim_accumulator = (0, 0, 0) print("Saving output to", save_dir) output = torch.squeeze(denormalise_lf(output)) cpu_output = np.around(output.cpu().detach().numpy()).astype(np.uint8) if (not args.no_eval) or args.get_diff: ground_truth = np.around( denormalise_lf(colour_images).numpy()).astype(np.uint8) grid_len = int(math.sqrt(grid_size)) for i in range(grid_size): row, col = i // grid_len, i % grid_len file_name = 'Colour{}{}.png'.format(row, col) save_location = os.path.join(save_dir, file_name) if i == 0: print("Saving images of size ", cpu_output[i].shape) image_warping.save_array_as_image(cpu_output[i], save_location) if args.get_diff: colour = ground_truth[i] diff = image_warping.get_diff_image(colour, cpu_output[i]) #diff = get_diff_image_floatint(res, colour) file_name = 'Diff{}{}.png'.format(row, col) save_location = os.path.join(save_dir, file_name) image_warping.save_array_as_image(diff, save_location) if not args.no_eval: img = ground_truth[i] file_name = 'GT_Colour{}{}.png'.format(row, col) save_location = os.path.join(save_dir, file_name) image_warping.save_array_as_image(img, save_location) psnr = evaluate.my_psnr(cpu_output[i], img) ssim = evaluate.ssim(cpu_output[i], img) psnr_accumulator = welford.update(psnr_accumulator, psnr) ssim_accumulator = welford.update(ssim_accumulator, ssim) psnr_mean, psnr_var, _ = welford.finalize(psnr_accumulator) ssim_mean, ssim_var, _ = welford.finalize(ssim_accumulator) print("For cnn, psnr average {:5f}, stddev {:5f}".format( psnr_mean, math.sqrt(psnr_var))) print("For cnn, ssim average {:5f}, stddev {:5f}".format( ssim_mean, math.sqrt(ssim_var)))
def do_one_demo(args, config, hdf5_file, model, sample_num, cuda): depth_grp = hdf5_file['val']['disparity'] # Create output directory if not args.no_save: base_dir = os.path.join(config['PATH']['output_dir'], 'warped') if not os.path.isdir(base_dir): pathlib.Path(base_dir).mkdir(parents=True, exist_ok=True) save_dir = get_sub_dir_for_saving(base_dir) model.eval() SNUM = sample_num start_time = time.time() print("Working on image", SNUM) depth_images = torch.squeeze( torch.tensor(depth_grp['images'][SNUM], dtype=torch.float32)) colour_grp = hdf5_file['val']['colour'] colour_images = torch.tensor(colour_grp['images'][SNUM], dtype=torch.float32) sample = { 'depth': depth_images, 'colour': colour_images, 'grid_size': depth_images.shape[0] } warped = data_transform.transform_to_warped(sample) warped = data_transform.stack(warped, 65) im_input = warped['inputs'].unsqueeze_(0) if cuda: im_input = im_input.cuda() output = model(im_input) im_input = im_input[:, :-3] output += im_input output = torch.clamp(output, 0.0, 1.0) time_taken = time.time() - start_time print("Time taken was {:.0f}s".format(time_taken)) grid_size = 64 psnr_accumulator = (0, 0, 0) ssim_accumulator = (0, 0, 0) if not args.no_save: print("Saving output to", save_dir) no_cnn_dir = os.path.join(save_dir, "no_cnn") cnn_dir = os.path.join(save_dir, "cnn") os.mkdir(cnn_dir) os.mkdir(no_cnn_dir) output = torch.squeeze(denormalise_lf(output)) output = data_transform.torch_unstack(output) im_input = im_input.cpu().detach() cpu_output = np.around(output.cpu().detach().numpy()).astype(np.uint8) if (not args.no_eval) or args.get_diff: ground_truth = np.around(denormalise_lf(colour_images).numpy()).astype( np.uint8) grid_len = int(math.sqrt(grid_size)) for i in range(grid_size): row, col = i // grid_len, i % grid_len if not args.no_save: file_name = 'Colour{}{}.png'.format(row, col) save_location = os.path.join(cnn_dir, file_name) if i == 0: print("Saving images of size ", cpu_output[i].shape) image_warping.save_array_as_image(cpu_output[i], save_location) if args.get_diff and not args.no_save: colour = ground_truth[i] diff = image_warping.get_diff_image(colour, cpu_output[i]) #diff = get_diff_image_floatint(res, colour) file_name = 'Diff{}{}.png'.format(row, col) save_location = os.path.join(cnn_dir, file_name) image_warping.save_array_as_image(diff, save_location) if not args.no_eval: img = ground_truth[i] if not args.no_save: file_name = 'GT_Colour{}{}.png'.format(row, col) save_location = os.path.join(save_dir, file_name) image_warping.save_array_as_image(img, save_location) psnr = evaluate.my_psnr(cpu_output[i], img) ssim = evaluate.ssim(cpu_output[i], img) psnr_accumulator = welford.update(psnr_accumulator, psnr) ssim_accumulator = welford.update(ssim_accumulator, ssim) psnr_mean, psnr_var, _ = welford.finalize(psnr_accumulator) ssim_mean, ssim_var, _ = welford.finalize(ssim_accumulator) print("For cnn, psnr average {:5f}, stddev {:5f}".format( psnr_mean, math.sqrt(psnr_var))) print("For cnn, ssim average {:5f}, stddev {:5f}".format( ssim_mean, math.sqrt(ssim_var))) psnr1 = psnr_mean ssim1 = ssim_mean psnr_accumulator = (0, 0, 0) ssim_accumulator = (0, 0, 0) psnr2, ssim2 = 0, 0 if args.no_cnn: squeeze_input = torch.squeeze(denormalise_lf(im_input)) squeeze_input = data_transform.torch_unstack(squeeze_input) cpu_input = np.around(squeeze_input.numpy()).astype(np.uint8) for i in range(grid_size): row, col = i // grid_len, i % grid_len if not args.no_save: file_name = 'Colour{}{}.png'.format(row, col) save_location = os.path.join(no_cnn_dir, file_name) if i == 0: print("Saving images of size ", cpu_input[i].shape) image_warping.save_array_as_image(cpu_input[i], save_location) if args.get_diff and not args.no_save: colour = ground_truth[i] diff = image_warping.get_diff_image(colour, cpu_output[i]) #diff = get_diff_image_floatint(res, colour) file_name = 'Diff{}{}.png'.format(row, col) save_location = os.path.join(no_cnn_dir, file_name) image_warping.save_array_as_image(diff, save_location) if not args.no_eval: img = ground_truth[i] psnr = evaluate.my_psnr(cpu_input[i], img) ssim = evaluate.ssim(cpu_input[i], img) psnr_accumulator = welford.update(psnr_accumulator, psnr) ssim_accumulator = welford.update(ssim_accumulator, ssim) psnr_mean, psnr_var, _ = welford.finalize(psnr_accumulator) ssim_mean, ssim_var, _ = welford.finalize(ssim_accumulator) print("For no cnn, psnr average {:5f}, stddev {:5f}".format( psnr_mean, math.sqrt(psnr_var))) print("For no cnn, ssim average {:5f}, stddev {:5f}".format( ssim_mean, math.sqrt(ssim_var))) psnr2, ssim2 = psnr_mean, ssim_mean return psnr1, ssim1, psnr2, ssim2
def main(args, config): hdf5_path = os.path.join(config['PATH']['output_dir'], config['PATH']['hdf5_name']) #warp_type = WARP_TYPE.TORCH_GPU warp_type = args.warp_type print("Performing image warping using {}".format(warp_type)) with h5py.File(hdf5_path, mode='r', libver='latest') as hdf5_file: grid_size = 64 grid_one_way = 8 sample_index = grid_size // 2 + (grid_one_way // 2) depth_grp = hdf5_file['val']['disparity'] overall_psnr_accum = (0, 0, 0) overall_ssim_accum = (0, 0, 0) for sample_num in range(args.nSamples): SNUM = sample_num print("Working on image", SNUM) depth_image = np.squeeze( depth_grp['images'][SNUM, sample_index, :192, :192]) #Hardcoded some values for now colour_grp = hdf5_file['val']['colour'] colour_image = colour_grp['images'][SNUM, sample_index] #Can later expand like 0000 if needed base_dir = os.path.join(config['PATH']['output_dir'], 'warped') get_diff = (config['DEFAULT']['should_get_diff'] == 'True') if not args.no_save: save_dir = get_sub_dir_for_saving(base_dir) print("Saving images to {}".format(save_dir)) else: print("Not saving images, only evaluating output") psnr_accumulator = (0, 0, 0) ssim_accumulator = (0, 0, 0) start_time = time() if warp_type == WARP_TYPE.TORCH_ALL: final = depth_rendering(colour_image, depth_image, lf_size=(64, 192, 192, 3)) print("Time taken was {:4f}".format(time() - start_time)) if warp_type == WARP_TYPE.TORCH_GPU: final = depth_rendering_gpu(colour_image, depth_image).cpu() print("Time taken was {:4f}".format(time() - start_time)) ref_pos = np.asarray([4, 4]) print("Reference position is ({}, {})".format(*ref_pos)) for i in range(8): for j in range(8): if warp_type == WARP_TYPE.FW: res = fw_warp_image(colour_image, depth_image, ref_pos, np.asarray([i, j])) elif warp_type == WARP_TYPE.SK: res = sk_warp(colour_image, depth_image, ref_pos, np.asarray([i, j]), preserve_range=True) elif warp_type == WARP_TYPE.SLOW: res = slow_fw_warp_image(colour_image, depth_image, ref_pos, np.asarray([i, j])) elif warp_type == WARP_TYPE.TORCH: res = np.around( torch_warp(colour_image, depth_image, ref_pos, np.asarray([i, j])).numpy()).astype( np.uint8) elif (warp_type == WARP_TYPE.TORCH_ALL or warp_type == WARP_TYPE.TORCH_GPU): res = np.around(final[i * 8 + j].numpy()).astype( np.uint8) if not args.no_save: file_name = 'Warped_Colour{}{}.png'.format(i, j) save_location = os.path.join(save_dir, file_name) save_array_as_image(res, save_location) idx = i * 8 + j file_name = 'GT_Colour{}{}.png'.format(i, j) save_location = os.path.join(save_dir, file_name) save_array_as_image(colour_grp['images'][SNUM][idx], save_location) if get_diff: colour = colour_grp['images'][SNUM, i * 8 + j] diff = get_diff_image(colour, res) #diff = get_diff_image_floatint(res, colour) file_name = 'Diff{}{}.png'.format(i, j) save_location = os.path.join(save_dir, file_name) save_array_as_image(diff, save_location) psnr = evaluate.my_psnr( res, colour_grp['images'][SNUM, i * 8 + j]) ssim = evaluate.ssim(res, colour_grp['images'][SNUM, i * 8 + j]) print("Position ({}, {}): PSNR {:4f}, SSIM {:4f}".format( i, j, psnr, ssim)) psnr_accumulator = welford.update(psnr_accumulator, psnr) ssim_accumulator = welford.update(ssim_accumulator, ssim) psnr_mean, psnr_var, _ = welford.finalize(psnr_accumulator) ssim_mean, ssim_var, _ = welford.finalize(ssim_accumulator) print("\npsnr average {:5f}, stddev {:5f}".format( psnr_mean, math.sqrt(psnr_var))) print("ssim average {:5f}, stddev {:5f}".format( ssim_mean, math.sqrt(ssim_var))) overall_psnr_accum = welford.update(overall_psnr_accum, psnr_mean) overall_ssim_accum = welford.update(overall_ssim_accum, ssim_mean) if args.nSamples > 1: psnr_mean, psnr_var, _ = welford.finalize(overall_psnr_accum) ssim_mean, ssim_var, _ = welford.finalize(overall_ssim_accum) print("\nOverall psnr average {:5f}, stddev {:5f}".format( psnr_mean, math.sqrt(psnr_var))) print("Overall ssim average {:5f}, stddev {:5f}".format( ssim_mean, math.sqrt(ssim_var)))
def save_looking_to_hdf5_group( sample_index, h5_canvas_list, camera, config): """Saves lf information to hdf5 group with images over columns then rows Keyword arguments: h5_canvas_list - tuples of groups and canvas numbers in a list for example (colour_group, 0) """ overall_start_time = time() cam = inviwopy.app.network.MeshClipping.camera accumulator_list = [] chs = config['channels'] # Gathers accumulators for welford mean and var for i, group_tuple in enumerate(h5_canvas_list): group = group_tuple[0] mean_shape = (1, ) + group["mean"].shape[1:] mean = np.zeros(mean_shape, np.float32) accumulator = (0, mean, 0) accumulator_list.append(accumulator) start_time = time() cam.lookFrom = camera[0] cam.lookTo = camera[1] cam.lookUp = camera[2] inviwo_utils.update() time_taken = time() - start_time sleep(1) for i, group_tuple in enumerate(h5_canvas_list): #assuming that there is two canvases group = group_tuple[0] canvas = ivw_helpers.get_canvas_by_id( inviwopy.app.network, group_tuple[1]) w_canvas = ivw_helpers.get_canvas_by_id( inviwopy.app.network, group_tuple[2]) assert (canvas.inputSize.customInputDimensions.value[0] // config["spatial_cols"] == group.attrs["lf_shape"][-1]), \ "canvas size x {} is not pixel dimension {} of h5".format( canvas.inputSize.customInputDimensions.value[0] // config["spatial_cols"], group.attrs["lf_shape"][-1] ) assert (canvas.inputSize.customInputDimensions.value[1] // config["spatial_rows"] == group.attrs["lf_shape"][-2]), \ "canvas size y {} is not pixel dimension {} of h5".format( canvas.inputSize.customInputDimensions.value[1] // config["spatial_rows"], group.attrs["lf_shape"][-2] ) total_im_data = ivw_helpers.get_image(canvas) group['camera_extrinsics'][sample_index] = \ [cam.lookFrom.x, cam.lookFrom.y, cam.lookFrom.z, cam.lookTo.x, cam.lookTo.y, cam.lookTo.z, cam.lookUp.x, cam.lookUp.y, cam.lookUp.z] # Inviwo stores data in a different indexing to regular # Store as descending y: C, H, W for idx in range(64): y_start = config["pixel_dim_y"] * (idx // 8) y_end = y_start + (config["pixel_dim_y"]) x_start = config["pixel_dim_x"] * (idx % 8) x_end = x_start + (config["pixel_dim_x"]) im_data = total_im_data[x_start:x_end, y_start:y_end] im_data = np.flipud(np.swapaxes(im_data, 0, 2)) im_data = im_data[::-1, ::-1, ...] group['images'][sample_index, idx] = im_data[:chs, ...] group['timing'][sample_index] = time_taken accumulator_list[i] = welford.update( accumulator_list[i], np.asarray(im_data[:chs, ...], dtype=np.float32)) total_im_data = ivw_helpers.get_image(w_canvas) # Inviwo stores data in a different indexing to regular # Store as descending y: C, H, W for idx in range(64): y_start = config["pixel_dim_y"] * (idx // 8) y_end = y_start + (config["pixel_dim_y"]) x_start = config["pixel_dim_x"] * (idx % 8) x_end = x_start + (config["pixel_dim_x"]) im_data = total_im_data[x_start:x_end, y_start:y_end] im_data = np.flipud(np.swapaxes(im_data, 0, 2)) im_data = im_data[::-1, ::-1, ...] group['warped'][sample_index, idx] = im_data[:chs, ...] for i, group_tuple in enumerate(h5_canvas_list): group = group_tuple[0] mean, var, _ = welford.finalize(accumulator_list[i]) group['mean'][sample_index, :, :, :] = mean group['var'][sample_index, :, :, :] = var print("Finished writing LF {0} in {1:.2f}".format( sample_index, time() - overall_start_time)) sleep(1)
def save_to_hdf5_group(self, sample_index, h5_canvas_list, config): """Saves lf information to hdf5 group with images over columns then rows Keyword arguments: h5_canvas_list - tuples of groups and canvas numbers in a list for example (colour_group, 0) """ overall_start_time = time() cam = inviwopy.app.network.MeshClipping.camera cam.lookUp = self.look_up accumulator_list = [] # Gathers accumulators for welford mean and var for i, group_tuple in enumerate(h5_canvas_list): group = group_tuple[0] mean_shape = (1, ) + group["mean"].shape[1:] mean = np.zeros(mean_shape, np.float32) accumulator = (0, mean, 0) accumulator_list.append(accumulator) for idx, val in enumerate(self.calculate_camera_array()): (look_from, look_to) = val start_time = time() cam.lookFrom = look_from cam.lookTo = look_to inviwo_utils.update() time_taken = time() - start_time #row_num, col_num = self.get_row_col_number(idx) for i, group_tuple in enumerate(h5_canvas_list): #assuming that there is one canvas group = group_tuple[0] canvas = inviwopy.app.network.canvases[group_tuple[1]] im_data = ivw_helpers.get_image(canvas) # Inviwo stores data in a different indexing to regular im_data = np.flipud(np.swapaxes(im_data, 0, 2)) im_data = im_data[::-1, ::-1, ...] assert (canvas.inputSize.customInputDimensions.value[0] == group.attrs["lf_shape"][-1]), \ "canvas size x {} is not pixel dimension {} of h5".format( canvas.inputSize.customInputDimensions.value[0], group.attrs["lf_shape"][-1] ) assert (canvas.inputSize.customInputDimensions.value[1] == group.attrs["lf_shape"][-2]), \ "canvas size y {} is not pixel dimension {} of h5".format( canvas.inputSize.customInputDimensions.value[1], group.attrs["lf_shape"][-2] ) group['images'][sample_index, idx] = im_data group['timing'][sample_index, idx] = time_taken group['camera_extrinsics'][sample_index, idx] = \ [cam.lookFrom.x, cam.lookFrom.y, cam.lookFrom.z, cam.lookTo.x, cam.lookTo.y, cam.lookTo.z, cam.lookUp.x, cam.lookUp.y, cam.lookUp.z] accumulator_list[i] = welford.update( accumulator_list[i], np.asarray(im_data, dtype=np.float32)) for i, group_tuple in enumerate(h5_canvas_list): group = group_tuple[0] mean, var, _ = welford.finalize(accumulator_list[i]) group['mean'][sample_index, :, :, :] = mean group['var'][sample_index, :, :, :] = var print("Finished writing LF {0} in {1:.2f} to {2}".format( sample_index, time() - overall_start_time, h5_canvas_list[0][0].name))
def main(config): hdf5_path = os.path.join(config['PATH']['output_dir'], config['PATH']['hdf5_name']) main_dir = config['PATH']['image_dir'] shared_metadata_keys = ('baseline', 'near', 'far', 'focal_length', 'fov') with h5py.File(hdf5_path, mode='w', libver='latest') as hdf5_file: hdf5_file.swmr_mode = True for set_type in ('train', 'val'): base_dir = os.path.join(main_dir, set_type) meta_dict = get_metadata(base_dir, config) sub_dirs = [ os.path.join(base_dir, el) for el in sorted(os.listdir(base_dir)) if os.path.isdir(os.path.join(base_dir, el)) ] hdf5_group = hdf5_file.create_group(set_type) num_samples = len(sub_dirs) #Handle the metadata meta_group = hdf5_group.create_group('metadata') for key in shared_metadata_keys: meta_group.create_dataset(key, (num_samples, ), np.float32) #Handle the depth group attributes depth = hdf5_group.create_group('disparity') # (num_images, grid_size, pixel_width, pixel_height, num_channels) depth.attrs['shape'] = [ num_samples, int(meta_dict['grid_rows']) * int(meta_dict['grid_cols']), int(meta_dict['pixels']), int(meta_dict['pixels']), 1 ] #Handle the colour group attributes colour = hdf5_group.create_group('colour') # Same shape as depth but with more channels temp = depth.attrs['shape'] temp[4] = 3 colour.attrs['shape'] = temp depth_image_shape = [ num_samples, int(meta_dict['pixels']), int(meta_dict['pixels']), 1 ] temp = np.copy(depth_image_shape) temp[3] = 3 colour_image_shape = temp #Save the images: dim = int(meta_dict['pixels']) depth.create_dataset('images', depth.attrs['shape'], np.float32, chunks=(1, 1, dim, dim, 1), compression="lzf", shuffle=True) depth.create_dataset('mean', depth_image_shape) depth.create_dataset('var', depth_image_shape) colour.create_dataset('images', colour.attrs['shape'], np.uint8, chunks=(1, 1, dim, dim, 3), compression="lzf", shuffle=True) colour.create_dataset('mean', colour_image_shape) colour.create_dataset('var', colour_image_shape) cols = int(meta_dict['grid_cols']) size = int(meta_dict['grid_rows']) * int(meta_dict['grid_cols']) for idx, loc in enumerate(sub_dirs): print(loc) depth_mean = np.zeros(depth_image_shape[1:-1], np.float32) colour_mean = np.zeros(colour_image_shape[1:], np.float32) depth_accumulator = (0, depth_mean, 0) colour_accumulator = (0, colour_mean, 0) meta = save_metadata(meta_group, idx, loc, shared_metadata_keys, config) for x in range(size): image_num = index_num_to_grid_loc(x, cols) depth_name = 'Depth' + image_num + '.npy' depth_loc = os.path.join(loc, depth_name) depth_data = np.load(depth_loc, fix_imports=False) depth_data = depth_to_disparity(depth_data, meta, dim) depth['images'][idx, x, :, :, 0] = depth_data depth_accumulator = (welford.update( depth_accumulator, depth_data)) colour_name = 'Colour' + image_num + '.png' colour_loc = os.path.join(loc, colour_name) colour_image = Image.open(colour_loc) colour_image.load() colour_data = np.asarray(colour_image, dtype=np.uint8) colour['images'][idx, x, :, :, :] = colour_data[:, :, :3] colour_accumulator = (welford.update( colour_accumulator, colour_data[:, :, :3])) (depth['mean'][idx, :, :, 0], depth['var'][idx, :, :, 0], _) = (welford.finalize(depth_accumulator)) (colour['mean'][idx, :, :, :], colour['var'][idx, :, :, :], _) = (welford.finalize(colour_accumulator)) print("Finished writing to", hdf5_path)