import cv2 class Detector(object): def __init__(self, model_arch, model_weights): self.model_arch = model_arch self.model_weights = model_weights def detect(self, image, thresh): raise NotImplementedError def crop(self, image, detections): crops = [] for det in detections: xmin = max(det[2], 0) ymin = max(det[3], 0) width = det[4] height = det[5] xmax = min(xmin+width, image.shape[1]) ymax = min(ymin+height, image.shape[0]) cut = image[ymin:ymax, xmin:xmax,:] crops.append(cut) return crops def draw(self, image, detections, im_scale=None): if im_scale is not None: image = cv2.resize(image, None, None, fx=im_scale, fy=im_scale, interpolation=cv2.INTER_LINEAR) detections = [[det[0],det[1],int(det[2]*im_scale),int(det[3]*im_scale),int(det[4]*im_scale),int(det[5]*im_scale)] for det in detections] for det in detections: xmin = det[2] ymin = det[3] width = det[4] height = det[5] xmax = xmin + width ymax = ymin + height cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2) return image