def __call__(self, data: TensorDict, rng=None): # Apply joint transforms if self.transform['joint'] is not None: num_train_images = len(data['train_images']) all_images = data['train_images'] + data['test_images'] all_images_trans = self.transform['joint'](*all_images) data['train_images'] = all_images_trans[:num_train_images] data['test_images'] = all_images_trans[num_train_images:] 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, rng) for a in data[s + '_anno'] ] # Crop image region centered at jittered_anno box try: crops, boxes = prutils.jittered_center_crop( data[s + '_images'], jittered_anno, data[s + '_anno'], self.search_area_factor[s], self.output_sz[s], scale_type=self.scale_type, border_type=self.border_type) except Exception as e: print('{}, anno: {}'.format(data['dataset'], data[s + '_anno'])) raise e # Apply transforms data[s + '_images'] = [self.transform[s](x) for x in crops] data[s + '_anno'] = boxes # Prepare output if self.mode == 'sequence': data = data.apply(prutils.stack_tensors) else: data = data.apply(lambda x: x[0] if isinstance(x, list) else x) return data
def __call__(self, data: TensorDict, rng=None): """ 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_iou' - """ # Apply joint transforms if self.transform['joint'] is not None: num_train_images = len(data['train_images']) all_images = data['train_images'] + data['test_images'] all_images_trans = self.transform['joint'](*all_images) data['train_images'] = all_images_trans[:num_train_images] data['test_images'] = all_images_trans[num_train_images:] 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, rng) for a in data[s + '_anno'] ] # Crop image region centered at jittered_anno box try: crops, boxes = prutils.jittered_center_crop( data[s + '_images'], jittered_anno, data[s + '_anno'], self.search_area_factor, self.output_sz) except Exception as e: print('{}, anno: {}'.format(data['dataset'], data[s + '_anno'])) raise e # Apply transforms data[s + '_images'] = [self.transform[s](x) for x in crops] data[s + '_anno'] = boxes # Generate proposals frame2_proposals, gt_iou = zip( *[self._generate_proposals(a, rng) for a in data['test_anno']]) data['test_proposals'] = list(frame2_proposals) data['proposal_iou'] = list(gt_iou) # Prepare output if self.mode == 'sequence': data = data.apply(prutils.stack_tensors) else: data = data.apply(lambda x: x[0] if isinstance(x, list) else x) return data
def __call__(self, data: TensorDict, rng=None): neg = data['neg'] # Apply joint transforms if self.transform['joint'] is not None: num_train_images = len(data['train_images']) all_images = data['train_images'] + data['test_images'] all_images_trans = self.transform['joint'](*all_images) data['train_images'] = all_images_trans[:num_train_images] data['test_images'] = all_images_trans[num_train_images:] 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, rng) for a in data[s + '_anno'] ] # Crop image region centered at jittered_anno box try: crops, boxes = prutils.jittered_center_crop( data[s + '_images'], jittered_anno, data[s + '_anno'], self.search_area_factor[s], self.output_sz[s], scale_type=self.scale_type, border_type=self.border_type) mask_crops, _ = prutils.jittered_center_crop( data[s + '_masks'], jittered_anno, data[s + '_anno'], self.search_area_factor[s], self.output_sz[s], scale_type=self.scale_type, border_type='zeropad') except Exception as e: print('{}, anno: {}'.format(data['dataset'], data[s + '_anno'])) raise e # Apply transforms data[s + '_images'] = [self.transform[s](x) for x in crops] data[s + '_anno'] = boxes data[s + '_masks'] = [ self.transform[s + '_mask'](x) for x in mask_crops ] # Prepare output if self.mode == 'sequence': data = data.apply(prutils.stack_tensors) else: data = data.apply(lambda x: x[0] if isinstance(x, list) else x) # Get labels if self.label_params is not None: assert data['test_anno'].shape[0] == 1 gt_box = data['test_anno'][0] gt_box[2:] += gt_box[:2] cls, delta, delta_weight, overlap = self._get_label(gt_box, neg) mask = data['test_masks'][0] if np.sum(mask) > 0: mask_weight = cls.max(axis=0, keepdims=True) else: mask_weight = np.zeros([1, cls.shape[1], cls.shape[2]], dtype=np.float32) mask = (mask > 0.5) * 2. - 1. data['label_cls'] = cls data['label_loc'] = delta data['label_loc_weight'] = delta_weight data['label_mask'] = mask data['label_mask_weight'] = mask_weight data.pop('train_anno') data.pop('test_anno') data.pop('train_masks') data.pop('test_masks') return data