| 113 | |
| 114 | |
| 115 | class MOTCenterNetDatasetMapper: |
| 116 | |
| 117 | def __init__(self, cfg, is_train=True): |
| 118 | if cfg.INPUT.CROP.ENABLED and is_train: |
| 119 | self.crop_gen = T.RandomCrop(cfg.INPUT.CROP.TYPE, cfg.INPUT.CROP.SIZE) |
| 120 | else: |
| 121 | self.crop_gen = None |
| 122 | |
| 123 | self.tfm_gens = build_MOT_transform_gen(cfg, is_train) |
| 124 | logging.getLogger(__name__).info( |
| 125 | "Full TransformGens used: {}, crop: {}".format(str(self.tfm_gens), str(self.crop_gen)) |
| 126 | ) |
| 127 | |
| 128 | self.img_format = cfg.INPUT.FORMAT |
| 129 | self.is_train = is_train |
| 130 | |
| 131 | def __call__(self, dataset_dicts): |
| 132 | if type(dataset_dicts).__name__ == 'tuple' and len(dataset_dicts) == 2: |
| 133 | dataset_dicts, gen_fake_img_flag = dataset_dicts[0], dataset_dicts[1] |
| 134 | dataset_dicts = copy.deepcopy(dataset_dicts) |
| 135 | crop_gen_on = True if self.crop_gen else False |
| 136 | for dataset_dict in dataset_dicts: |
| 137 | image = utils.read_image(dataset_dict["file_name"], format=self.img_format) |
| 138 | utils.check_image_size(dataset_dict, image) |
| 139 | |
| 140 | if "annotations" not in dataset_dict: |
| 141 | image, transforms = T.apply_transform_gens( |
| 142 | ([self.crop_gen] if self.crop_gen else []) + self.tfm_gens, image |
| 143 | ) |
| 144 | else: |
| 145 | if crop_gen_on: |
| 146 | crop_tfm = utils.gen_crop_transform_with_instance( |
| 147 | self.crop_gen.get_crop_size(image.shape[:2]), |
| 148 | image.shape[:2], |
| 149 | np.random.choice(dataset_dict["annotations"]), |
| 150 | ) |
| 151 | image = crop_tfm.apply_image(image) |
| 152 | image, transforms = T.apply_transform_gens(self.tfm_gens, image) |
| 153 | if crop_gen_on: |
| 154 | transforms = crop_tfm + transforms |
| 155 | |
| 156 | image_shape = image.shape[:2] # h, w |
| 157 | dataset_dict["image"] = torch.as_tensor(np.ascontiguousarray(image.transpose(2, 0, 1))) |
| 158 | |
| 159 | if not self.is_train: |
| 160 | dataset_dict.pop("annotations", None) |
| 161 | else: |
| 162 | if "annotations" in dataset_dict: |
| 163 | annos = [] |
| 164 | for obj in dataset_dict.pop("annotations"): |
| 165 | annos.append(transform_instance_annotations(obj, transforms, image_shape)) |
| 166 | instances = annotations_to_instances(annos, image_shape) |
| 167 | dataset_dict["instances"] = utils.filter_empty_instances(instances) |
| 168 | |
| 169 | if len(dataset_dicts) == 1: |
| 170 | return dataset_dicts[0] |
| 171 | else: |
| 172 | return dataset_dicts |
no outgoing calls
no test coverage detected