def test_HWC2CHW_comp(plot=False): """ Test HWC2CHW between python and c image augmentation """ logger.info("Test HWC2CHW with c_transform and py_transform comparison") # First dataset data1 = ds.TFRecordDataset(DATA_DIR, SCHEMA_DIR, columns_list=["image"], shuffle=False) decode_op = c_vision.Decode() hwc2chw_op = c_vision.HWC2CHW() data1 = data1.map(input_columns=["image"], operations=decode_op) data1 = data1.map(input_columns=["image"], operations=hwc2chw_op) # Second dataset data2 = ds.TFRecordDataset(DATA_DIR, SCHEMA_DIR, columns_list=["image"], shuffle=False) transforms = [ py_vision.Decode(), py_vision.ToTensor(), py_vision.HWC2CHW() ] transform = py_vision.ComposeOp(transforms) data2 = data2.map(input_columns=["image"], operations=transform()) image_c_transposed = [] image_py_transposed = [] for item1, item2 in zip(data1.create_dict_iterator(), data2.create_dict_iterator()): c_image = item1["image"] py_image = (item2["image"].transpose(1, 2, 0) * 255).astype(np.uint8) # compare images between that applying c_transform and py_transform mse = diff_mse(py_image, c_image) # the images aren't exactly the same due to rounding error assert mse < 0.001 image_c_transposed.append(item1["image"].copy()) image_py_transposed.append(item2["image"].copy()) if plot: visualize(image_c_transposed, image_py_transposed)
def create_yolo_dataset(mindrecord_dir, batch_size=32, repeat_num=10, device_num=1, rank=0, is_training=True, num_parallel_workers=8): """Creatr YOLOv3 dataset with MindDataset.""" ds = de.MindDataset(mindrecord_dir, columns_list=["image", "annotation"], num_shards=device_num, shard_id=rank, num_parallel_workers=num_parallel_workers, shuffle=is_training) decode = C.Decode() ds = ds.map(input_columns=["image"], operations=decode) compose_map_func = (lambda image, annotation: preprocess_fn( image, annotation, is_training)) if is_training: hwc_to_chw = P.HWC2CHW() ds = ds.map(input_columns=["image", "annotation"], output_columns=[ "image", "bbox_1", "bbox_2", "bbox_3", "gt_box1", "gt_box2", "gt_box3" ], columns_order=[ "image", "bbox_1", "bbox_2", "bbox_3", "gt_box1", "gt_box2", "gt_box3" ], operations=compose_map_func, num_parallel_workers=num_parallel_workers) ds = ds.map(input_columns=["image"], operations=hwc_to_chw, num_parallel_workers=num_parallel_workers) ds = ds.shuffle(buffer_size=256) ds = ds.batch(batch_size, drop_remainder=True) ds = ds.repeat(repeat_num) else: ds = ds.map(input_columns=["image", "annotation"], output_columns=["image", "image_shape", "annotation"], columns_order=["image", "image_shape", "annotation"], operations=compose_map_func, num_parallel_workers=num_parallel_workers) return ds
def create_yolo_dataset(image_dir, anno_path, batch_size=32, repeat_num=10, device_num=1, rank=0, is_training=True, num_parallel_workers=8): """Creatr YOLOv3 dataset with GeneratorDataset.""" yolo_dataset = YoloDataset(image_dir=image_dir, anno_path=anno_path) distributed_sampler = DistributedSampler(yolo_dataset.dataset_size, batch_size, device_num, rank) ds = de.GeneratorDataset(yolo_dataset, column_names=["image", "annotation"], sampler=distributed_sampler) ds.set_dataset_size(len(distributed_sampler)) compose_map_func = (lambda image, annotation: preprocess_fn(image, annotation, is_training)) hwc_to_chw = P.HWC2CHW() ds = ds.map(input_columns=["image", "annotation"], output_columns=["image", "bbox_1", "bbox_2", "bbox_3", "gt_box1", "gt_box2", "gt_box3"], columns_order=["image", "bbox_1", "bbox_2", "bbox_3", "gt_box1", "gt_box2", "gt_box3"], operations=compose_map_func, num_parallel_workers=num_parallel_workers) ds = ds.map(input_columns=["image"], operations=hwc_to_chw, num_parallel_workers=num_parallel_workers) ds = ds.shuffle(buffer_size=256) ds = ds.batch(batch_size, drop_remainder=True) ds = ds.repeat(repeat_num) return ds