目标检测模型代码

def letterbox(img, new_shape=(640, 640), color=(114, 114, 114), auto=False, scaleFill=False, scaleup=True):
    # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232
    shape = img.shape[:2]  # current shape [height, width]
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)
 
    # Scale ratio (new / old)
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  # only scale down, do not scale up (for better test mAP)
        r = min(r, 1.0)
 
    # Compute padding
    ratio = r, r  # width, height ratios
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
    if auto:  # minimum rectangle
        dw, dh = np.mod(dw, 64), np.mod(dh, 64)  # wh padding
    elif scaleFill:  # stretch
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios
 
    dw /= 2  # divide padding into 2 sides
    dh /= 2
 
    if shape[::-1] != new_unpad:  # resize
        img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
    return img, ratio, (dw, dh)

class YoloV5(Model):
    def __init__(self, model_path, acl_init=True):
        super().__init__(model_path, acl_init)
        self.neth = 640
        self.netw = 640
        self.conf_threshold = 0.1
        dic = {0: 'left',
               1: 'right',
               2: 'stop',
               3: 'turnaround'}
        self.names = ['person', 'sports_ball', 'bicycle', 'motorcycle', 'car', 'bus', 'truck'] * 12
        self.object_list = ['person', 'sports_ball', 'bicycle', 'motorcycle', 'car', 'bus', 'truck']
        self.names = list(dic.values())
        self.object_list = list(dic.values())
 
    def infer(self, img_bgr):
        imgh, imgw = img_bgr.shape[0], img_bgr.shape[1]
        imginfo = np.array([self.neth, self.netw, imgh, imgw], dtype=np.float16)
        img_padding = letterbox(img_bgr, new_shape=(self.neth, self.netw))[0]  # padding resize bgr
 
        img = []
 
        img.append(img_padding)
        img = np.stack(img, axis=0)
        img = img[..., ::-1].transpose(0, 3, 1, 2)  # BGR tp RGB
        image_np = np.array(img, dtype=np.float32)
        image_np_expanded = image_np / 255.0
        img = np.ascontiguousarray(image_np_expanded).astype(np.float16) #将tensor的内存连续排列
        result = self.execute([img, imginfo]) #调用推理接口
        batch_boxout, boxnum = result
 
        pred_boxes = []
        idx = 0
        num_det = int(boxnum[idx][0])
        bbox = batch_boxout[idx][:num_det * 6].reshape(6, -1).transpose().astype(np.float32)  # 6xN -> Nx6
 
        for idx, class_id in enumerate(bbox[:, 5]):
            obj_name = self.names[int(bbox[idx][5])]
            if not obj_name in self.object_list:
                continue
            confidence = bbox[idx][4]
            if float(confidence) < self.conf_threshold:
                continue
            x1 = int(bbox[idx][0])
            y1 = int(bbox[idx][1])
            x2 = int(bbox[idx][2])
            y2 = int(bbox[idx][3])
 
            pred_boxes.append([x1, y1, x2, y2, obj_name, confidence]) #获取推理结果
 
        return pred_boxes