| 34 | |
| 35 | |
| 36 | class TrainPre(object): |
| 37 | def __init__(self, norm_mean, norm_std, sign=False, config=None): |
| 38 | self.config = config |
| 39 | self.norm_mean = norm_mean |
| 40 | self.norm_std = norm_std |
| 41 | self.sign = sign |
| 42 | |
| 43 | def __call__(self, rgb, gt, modal_x): |
| 44 | rgb, gt, modal_x = random_mirror(rgb, gt, modal_x) |
| 45 | if self.config.train_scale_array is not None: |
| 46 | rgb, gt, modal_x, scale = random_scale(rgb, gt, modal_x, self.config.train_scale_array) |
| 47 | |
| 48 | rgb = normalize(rgb, self.norm_mean, self.norm_std) |
| 49 | if self.sign: |
| 50 | modal_x = normalize(modal_x, [0.48, 0.48, 0.48], [0.28, 0.28, 0.28]) # [0.5,0.5,0.5] |
| 51 | else: |
| 52 | modal_x = normalize(modal_x, self.norm_mean, self.norm_std) |
| 53 | |
| 54 | # return rgb.transpose(2, 0, 1), gt, modal_x.transpose(2, 0, 1) |
| 55 | |
| 56 | crop_size = (self.config.image_height, self.config.image_width) |
| 57 | crop_pos = generate_random_crop_pos(rgb.shape[:2], crop_size) |
| 58 | |
| 59 | p_rgb, _ = random_crop_pad_to_shape(rgb, crop_pos, crop_size, 0) |
| 60 | p_gt, _ = random_crop_pad_to_shape(gt, crop_pos, crop_size, 255) |
| 61 | p_modal_x, _ = random_crop_pad_to_shape(modal_x, crop_pos, crop_size, 0) |
| 62 | |
| 63 | p_rgb = p_rgb.transpose(2, 0, 1) |
| 64 | p_modal_x = p_modal_x.transpose(2, 0, 1) |
| 65 | # p_rgb = p_rgb |
| 66 | # p_modal_x = p_modal_x |
| 67 | |
| 68 | return p_rgb, p_gt, p_modal_x |
| 69 | |
| 70 | |
| 71 | def seed_worker(worker_id): |