def val_preproc(sample, preproc_args):
ref, ref_cam, srcs, srcs_cam, gt, masks, skip = [
sample[attr] for attr in
['ref', 'ref_cam', 'srcs', 'srcs_cam', 'gt', 'masks', 'skip']
]
ref, *srcs = [center_image(img) for img in [ref] + srcs]
ref, ref_cam, srcs, srcs_cam, gt, masks = resize(
[ref, ref_cam, srcs, srcs_cam, gt, masks],
preproc_args['resize_width'], preproc_args['resize_height'])
ref, ref_cam, srcs, srcs_cam, gt, masks = center_crop(
[ref, ref_cam, srcs, srcs_cam, gt, masks], preproc_args['crop_width'],
preproc_args['crop_height'])
ref, *srcs, gt = to_channel_first([ref] + srcs + [gt])
masks = to_channel_first(masks)
srcs, srcs_cam, masks = [
np.stack(arr_list, axis=0) for arr_list in [srcs, srcs_cam, masks]
]
return {
'ref': ref, # 3hw
'ref_cam': ref_cam, # 244
'srcs': srcs, # v3hw
'srcs_cam': srcs_cam, # v244
'gt': gt, # 1hw
'masks': masks, # v1hw
'skip': skip # scalar
}
def train_preproc(sample, preproc_args):
ref, ref_cam, srcs, srcs_cam, gt, masks, skip = [
sample[attr] for attr in
['ref', 'ref_cam', 'srcs', 'srcs_cam', 'gt', 'masks', 'skip']
]
ref, *srcs = [
random_contrast(img, strength_range=[0.3, 1.5]) for img in [ref] + srcs
]
ref, *srcs = [
random_brightness(img, max_abs_change=50) for img in [ref] + srcs
]
ref, *srcs = [motion_blur(img, max_kernel_size=3) for img in [ref] + srcs]
ref, *srcs = [center_image(img) for img in [ref] + srcs]
ref, ref_cam, srcs, srcs_cam, gt, masks = resize(
[ref, ref_cam, srcs, srcs_cam, gt, masks],
preproc_args['resize_width'], preproc_args['resize_height'])
ref, ref_cam, srcs, srcs_cam, gt, masks = center_crop(
[ref, ref_cam, srcs, srcs_cam, gt, masks], preproc_args['crop_width'],
preproc_args['crop_height'])
ref, *srcs, gt = to_channel_first([ref] + srcs + [gt])
masks = to_channel_first(masks)
srcs, srcs_cam, masks = [
np.stack(arr_list, axis=0) for arr_list in [srcs, srcs_cam, masks]
]
return {
'ref': ref, # 3hw
'ref_cam': ref_cam, # 244
'srcs': srcs, # v3hw
'srcs_cam': srcs_cam, # v244
'gt': gt, # 1hw
'masks': masks, # v1hw
'skip': skip # scalar
}