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