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 }