'''Define classes and functions for interfacing with SageMaker Ground Truth object detection. ''' import os import imageio import matplotlib.pyplot as plt import numpy as np class BoundingBox: '''Bounding box for an object in an image.''' def __init__(self, image_id=None, boxdata=None): self.image_id = image_id if boxdata: for datum in boxdata: setattr(self, datum, boxdata[datum]) def __repr__(self): return 'Box for image {}'.format(self.image_id) def compute_bb_data(self): '''Compute the parameters used for IoU.''' image = self.image self.xmin = self.left/image.width self.xmax = (self.left + self.width)/image.width self.ymin = self.top/image.height self.ymax = (self.top + self.height)/image.height class WorkerBoundingBox(BoundingBox): '''Bounding box for an object in an image produced by a worker.''' def __init__(self, image_id=None, worker_id=None, boxdata=None): self.worker_id = worker_id super().__init__(image_id=image_id, boxdata=boxdata) class GroundTruthBox(BoundingBox): '''Bounding box for an object in an image produced by a worker.''' def __init__(self, image_id=None, oiddata=None, image=None): self.image = image self.class_name = oiddata[0] xmin, xmax, ymin, ymax = [float(datum) for datum in oiddata[1:]] self.xmin = xmin self.ymin = ymin self.xmax = xmax self.ymax = ymax imw = image.width imh = image.height boxdata = {'height': (ymax-ymin)*imh, 'width': (xmax-xmin)*imw, 'left': xmin*imw, 'top': ymin*imh} super().__init__(image_id=image_id, boxdata=boxdata) class BoxedImage: '''Image with bounding boxes.''' def __init__(self, id=None, consolidated_boxes=None, worker_boxes=None, gt_boxes=None, uri=None, size=None): self.id = id self.uri = uri if uri: self.filename = uri.split('/')[-1] self.oid_id = self.filename.split('.')[0] else: self.filename = None self.oid_id = None self.local = None self.im = None if size: self.width = size['width'] self.depth = size['depth'] self.height = size['height'] self.shape = self.width, self.height, self.depth if consolidated_boxes: self.consolidated_boxes = consolidated_boxes else: self.consolidated_boxes = [] if worker_boxes: self.worker_boxes = worker_boxes else: self.worker_boxes = [] if gt_boxes: self.gt_boxes = gt_boxes else: self.gt_boxes = [] def __repr__(self): return 'Image{}'.format(self.id) def n_consolidated_boxes(self): '''Count the number of consolidated boxes.''' return len(self.consolidated_boxes) def n_worker_boxes(self): return len(self.worker_boxes) def download(self, directory): target_fname = os.path.join( directory, self.uri.split('/')[-1]) if not os.path.isfile(target_fname): os.system(f'aws s3 cp {self.uri} {target_fname}') self.local = target_fname def imread(self): '''Cache the image reading process.''' try: return imageio.imread(self.local) except OSError: print("You need to download this image first. " "Use this_image.download(local_directory).") raise def plot_bbs(self, ax, bbs, img_kwargs, box_kwargs, **kwargs): '''Master function for plotting images with bounding boxes.''' img = self.imread() ax.imshow(img, **img_kwargs) imh, imw, *_ = img.shape box_kwargs['fill'] = None if kwargs.get('worker', False): # Give each worker a color. worker_colors = {} worker_count = 0 for bb in bbs: worker = bb.worker_id if worker not in worker_colors: worker_colors[worker] = 'C' + str((9-worker_count) % 10) worker_count += 1 rec = plt.Rectangle((bb.left, bb.top), bb.width, bb.height, edgecolor=worker_colors[worker], **box_kwargs) ax.add_patch(rec) else: for bb in bbs: rec = plt.Rectangle( (bb.left, bb.top), bb.width, bb.height, **box_kwargs) ax.add_patch(rec) ax.axis('off') def plot_consolidated_bbs(self, ax, img_kwargs={}, box_kwargs={'edgecolor': 'blue', 'lw': 3}): '''Plot the consolidated boxes.''' self.plot_bbs(ax, self.consolidated_boxes, img_kwargs=img_kwargs, box_kwargs=box_kwargs) def plot_worker_bbs(self, ax, img_kwargs={}, box_kwargs={'lw': 2}): '''Plot the individual worker boxes.''' self.plot_bbs(ax, self.worker_boxes, worker=True, img_kwargs=img_kwargs, box_kwargs=box_kwargs) def plot_gt_bbs(self, ax, img_kwargs={}, box_kwargs={'edgecolor': 'lime', 'lw': 3}): '''Plot the ground truth (Open Image Dataset) boxes.''' self.plot_bbs(ax, self.gt_boxes, img_kwargs=img_kwargs, box_kwargs=box_kwargs) def compute_img_confidence(self): ''' Compute the mean bb confidence. ''' if len(self.consolidated_boxes) > 0: return np.mean([box.confidence for box in self.consolidated_boxes]) else: return 0 def compute_iou_bb(self): '''Compute the mean intersection over union for a collection of bounding boxes. ''' # Precompute data for the consolidated boxes if necessary. for box in self.consolidated_boxes: try: box.xmin except AttributeError: box.compute_bb_data() # Make the numpy arrays. if self.gt_boxes: gts = np.vstack([(box.xmin, box.ymin, box.xmax, box.ymax) for box in self.gt_boxes]) else: gts = [] if self.consolidated_boxes: preds = np.vstack([(box.xmin, box.ymin, box.xmax, box.ymax) for box in self.consolidated_boxes]) else: preds = [] confs = np.array([box.confidence for box in self.consolidated_boxes]) if len(preds) == 0 and len(gts) == 0: return 1. if len(preds) == 0 or len(gts) == 0: return 0. preds = preds[np.argsort(confs.flatten())][::-1] is_pred_assigned_to_gt = [False] * len(gts) pred_areas = (preds[:, 2] - preds[:, 0]) * \ (preds[:, 3] - preds[:, 1]) gt_areas = (gts[:, 2] - gts[:, 0]) * (gts[:, 3] - gts[:, 1]) all_ious = [] for pred_id, pred in enumerate(preds): best_iou = 0 best_id = -1 for gt_id, gt in enumerate(gts): if is_pred_assigned_to_gt[gt_id]: continue x1 = max(gt[0], pred[0]) y1 = max(gt[1], pred[1]) x2 = min(gt[2], pred[2]) y2 = min(gt[3], pred[3]) iw = max(0, x2 - x1) ih = max(0, y2 - y1) inter = iw * ih iou = inter / \ (pred_areas[pred_id] + gt_areas[gt_id] - inter) if iou > best_iou: best_iou = iou best_id = gt_id if best_id != -1: is_pred_assigned_to_gt[best_id] = True # True positive! Store the IoU. all_ious.append(best_iou) else: # 0 IoU for each unmatched gt (false-negative). all_ious.append(0.) # 0 IoU for each unmatched prediction (false-positive). all_ious.extend([0.] * (len(is_pred_assigned_to_gt) - sum(is_pred_assigned_to_gt))) return np.mean(all_ious) def group_miou(imgs): '''Compute the mIoU for a group of images. Args: imgs: list of BoxedImages, with consolidated_boxes and gt_boxes. Returns: mIoU calculated over the bounding boxes in the group. ''' # Create a notional BoxedImage with bounding boxes from imgs. all_consolidated_boxes = [box for img in imgs for box in img.consolidated_boxes] all_gt_boxes = [box for img in imgs for box in img.gt_boxes] notional_image = BoxedImage(consolidated_boxes=all_consolidated_boxes, gt_boxes=all_gt_boxes) # Compute and return the mIoU. return notional_image.compute_iou_bb()