def __call__(self, sample, opts):
image, bb = sample['image'], sample['bb']
image_x2 = None
img_x2 = None
if ('image_x2' in sample):
image_x2 = sample['image_x2']
h, w = image.shape[:2]
if isinstance(self.output_size, int):
if h > w:
new_h, new_w = self.output_size * h / w, self.output_size
else:
new_h, new_w = self.output_size, self.output_size * w / h
else:
new_h, new_w = self.output_size
new_h, new_w = int(new_h), int(new_w)
# make sure that gray image has 3 channels
img = cv2.resize(image, (new_h, new_w), interpolation=cv2.INTER_CUBIC)
if image_x2 is not None:
img_x2 = cv2.resize(image_x2, (new_h * 2, new_w * 2),
interpolation=cv2.INTER_CUBIC)
bbox = BoundingBox(bb[0], bb[1], bb[2], bb[3])
bbox.scale(opts['search_region'])
return {'image': img, 'image_x2': img_x2, 'bb': bbox.get_bb_list()}
def crop_sample(sample):
"""
Given a sample image with bounding box, this method returns the image crop
at the bounding box location with twice the width and height for context.
"""
output_sample = {}
opts = {}
image, bb = sample['image'], sample['bb']
orig_bbox = BoundingBox(bb[0], bb[1], bb[2], bb[3])
(output_image, pad_image_location,
edge_spacing_x, edge_spacing_y) = cropPadImage(orig_bbox, image)
new_bbox = BoundingBox(0, 0, 0, 0)
new_bbox = new_bbox.recenter(pad_image_location,
edge_spacing_x,
edge_spacing_y,
new_bbox)
output_sample['image'] = output_image
output_sample['bb'] = new_bbox.get_bb_list()
# additional options for visualization
opts['edge_spacing_x'] = edge_spacing_x
opts['edge_spacing_y'] = edge_spacing_y
opts['search_location'] = pad_image_location
opts['search_region'] = output_image
return output_sample, opts
def shift_crop_training_sample(sample, bb_params):
"""
Given an image with bounding box, this method randomly shifts the box and
generates a training example. It returns current image crop with shifted
box (with respect to current image).
"""
output_sample = {}
opts = {}
currimg = sample['image']
currbb = sample['bb']
bbox_curr_gt = BoundingBox(currbb[0], currbb[1], currbb[2], currbb[3])
bbox_curr_shift = BoundingBox(0, 0, 0, 0)
bbox_curr_shift = bbox_curr_gt.shift(currimg,
bb_params['lambda_scale_frac'],
bb_params['lambda_shift_frac'],
bb_params['min_scale'],
bb_params['max_scale'], True,
bbox_curr_shift)
(rand_search_region, rand_search_location, edge_spacing_x,
edge_spacing_y) = cropPadImage(bbox_curr_shift, currimg)
bbox_curr_gt = BoundingBox(currbb[0], currbb[1], currbb[2], currbb[3])
bbox_gt_recentered = BoundingBox(0, 0, 0, 0)
bbox_gt_recentered = bbox_curr_gt.recenter(rand_search_location,
edge_spacing_x, edge_spacing_y,
bbox_gt_recentered)
output_sample['image'] = rand_search_region
output_sample['bb'] = bbox_gt_recentered.get_bb_list()
# additional options for visualization
opts['edge_spacing_x'] = edge_spacing_x
opts['edge_spacing_y'] = edge_spacing_y
opts['search_location'] = rand_search_location
opts['search_region'] = rand_search_region
return output_sample, opts
def get_rect(self, sample):
"""
Regresses the bounding box coordinates in the original image dimensions
for an input sample.
"""
x1, x2 = sample['previmg'], sample['currimg']
x1 = x1.unsqueeze(0).to(self.device)
x2 = x2.unsqueeze(0).to(self.device)
y = self.model(x1, x2)
bb = y.data.cpu().numpy().transpose((1, 0))
bb = bb[:, 0]
bbox = BoundingBox(bb[0], bb[1], bb[2], bb[3])
# inplace conversion
bbox.unscale(self.opts['search_region'])
bbox.uncenter(self.curr_img, self.opts['search_location'],
self.opts['edge_spacing_x'], self.opts['edge_spacing_y'])
return bbox.get_bb_list()