def __call__(self, data: TensorDict): # Apply joint transformations. i.e. All train/test frames in a sequence are applied the transformation with the # same parameters if self.transform['joint'] is not None: data['train_images'], data['train_anno'], data['train_masks'] = self.transform['joint']( image=data['train_images'], bbox=data['train_anno'], mask=data['train_masks']) data['test_images'], data['test_anno'], data['test_masks'] = self.transform['joint']( image=data['test_images'], bbox=data['test_anno'], mask=data['test_masks'], new_roll=self.new_roll) for s in ['train', 'test']: assert self.mode == 'sequence' or len(data[s + '_images']) == 1, \ "In pair mode, num train/test frames must be 1" # Add a uniform noise to the center pos jittered_anno = [self._get_jittered_box(a, s) for a in data[s + '_anno']] orig_anno = data[s + '_anno'] # Extract a crop containing the target crops, boxes, mask_crops = prutils.target_image_crop(data[s + '_images'], jittered_anno, data[s + '_anno'], self.search_area_factor, self.output_sz, mode=self.crop_type, max_scale_change=self.max_scale_change, masks=data[s + '_masks']) # Apply independent transformations to each image data[s + '_images'], data[s + '_anno'], data[s + '_masks'] = self.transform[s](image=crops, bbox=boxes, mask=mask_crops, joint=False) # Prepare output if self.mode == 'sequence': data = data.apply(stack_tensors) else: data = data.apply(lambda x: x[0] if isinstance(x, list) else x) return data
def __call__(self, data: TensorDict): """ args: data - The input data, should contain the following fields: 'train_images', test_images', 'train_anno', 'test_anno' returns: TensorDict - output data block with following fields: 'train_images', 'test_images', 'train_anno', 'test_anno', 'test_proposals', 'proposal_density', 'gt_density', 'test_label' (optional), 'train_label' (optional), 'test_label_density' (optional), 'train_label_density' (optional) """ if self.transform['joint'] is not None: data['train_images'], data['train_anno'] = self.transform['joint'](image=data['train_images'], bbox=data['train_anno']) data['test_images'], data['test_anno'] = self.transform['joint'](image=data['test_images'], bbox=data['test_anno'], new_roll=False) for s in ['train', 'test']: assert self.mode == 'sequence' or len(data[s + '_images']) == 1, \ "In pair mode, num train/test frames must be 1" # Add a uniform noise to the center pos jittered_anno = [self._get_jittered_box(a, s) for a in data[s + '_anno']] crops, boxes = prutils.target_image_crop(data[s + '_images'], jittered_anno, data[s + '_anno'], self.search_area_factor, self.output_sz, mode=self.crop_type, max_scale_change=self.max_scale_change) data[s + '_images'], data[s + '_anno'] = self.transform[s](image=crops, bbox=boxes, joint=False) # Generate proposals proposals, proposal_density, gt_density = zip(*[self._generate_proposals(a) for a in data['test_anno']]) data['test_proposals'] = proposals data['proposal_density'] = proposal_density data['gt_density'] = gt_density for s in ['train', 'test']: is_distractor = data.get('is_distractor_{}_frame'.format(s), None) if is_distractor is not None: for is_dist, box in zip(is_distractor, data[s+'_anno']): if is_dist: box[0] = 99999999.9 box[1] = 99999999.9 # Prepare output if self.mode == 'sequence': data = data.apply(stack_tensors) else: data = data.apply(lambda x: x[0] if isinstance(x, list) else x) # Generate label functions if self.label_function_params is not None: data['train_label'] = self._generate_label_function(data['train_anno']) data['test_label'] = self._generate_label_function(data['test_anno']) if self.label_density_params is not None: data['train_label_density'] = self._generate_label_density(data['train_anno']) data['test_label_density'] = self._generate_label_density(data['test_anno']) return data