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