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 val_preproc(sample, preproc_args):
ref, ref_cam, srcs, srcs_cam, ref_depth, srcs_depth, ref_probs = [
sample[attr] for attr in [
'ref', 'ref_cam', 'srcs', 'srcs_cam', 'ref_depth', 'srcs_depth',
'ref_probs'
]
]
ref, *srcs = to_channel_first([ref] + srcs)
ref_depth, *srcs_depth = to_channel_first([ref_depth] + srcs_depth)
ref_probs = to_channel_first(ref_probs)
srcs, srcs_cam, srcs_depth, ref_probs = [
np.stack(arr_list, axis=0)
for arr_list in [srcs, srcs_cam, srcs_depth, ref_probs]
]
return {
'id': sample['id'],
'ref': ref, # 3hw
'ref_cam': ref_cam, # 244
'ref_depth': ref_depth, # 1hw
'srcs': srcs, # v3hw
'srcs_cam': srcs_cam, # v244
'srcs_depth': srcs_depth, # v1hw
'ref_probs': ref_probs, # 31hw
'skip': sample['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
}
def channel_first_wrapper(sample):
ref, ref_cam, srcs, srcs_cam, gt, masks = sample
ref, *srcs = to_channel_first([ref] + srcs)
return ref, ref_cam, srcs, srcs_cam, gt, masks
