third_party/PIPNet/FaceBoxesV2/detector.py

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