YOLOv5模型的API调用方式

调用API方式1

1. 首先同目录下需要复制YOLOv5代码的models文件、utils文件、export文件。
2. 返回格式为列表形式，存储所有的目标位置，置信度等

def detctapi(im):
    # Load Yolov5 model
    model_path = r'weights/yolo.pt'
    data = r'data/coco128.yaml'
    device = torch.device("cuda:0")

    @torch.no_grad()
    def init():
        weights = model_path
        device = 'cuda:0'  # cuda device, i.e. 0 or 0,1,2,3 or

        half = True  # use FP16 half-precision inference
        device = select_device(device)
        w = str(weights[0] if isinstance(weights, list) else weights)
        # model = torch.jit.load(w) if 'torchscript' in w else attempt_load(weights, device=device)
        model = DetectMultiBackend(w, device=device, dnn=False, data=data, fp16=half)

        model.eval()
        return model

    # model = torch.load(r'D:\yanyi\xianyu\1\facedetector_facemask3\facedetector_facemask3\untitled\weights\yolo.pt', map_location=torch.device('cuda'))
    model = init()
    conf_thres = 0.25  # confidence threshold
    iou_thres = 0.45  # NMS IOU threshold
    classes = int(0)
    max_det = 1000  # maximum detections per image


    # Load input image
    stride = 32
    imgsz = [320, 320]
    imgsz = check_img_size(imgsz, s=stride)  # check image size
    # Process image here
    fake_result = {}
    fake_result["model_data"] = {"objects": []}

    img = letterbox(im, imgsz)[0]
    print(img.shape)
    img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
    img = np.ascontiguousarray(img)
    img = torch.from_numpy(img).to(device)
    img = img.half() if model.fp16 else img.float()  # uint8 to fp16/32
    img /= 255  # 0 - 255 to 0.0 - 1.0
    if len(img.shape) == 3:
        img = img[None]  # expand for batch dim
    names = {
        0: 'Right_mask',
        1: 'Error_mask',
        2: 'No_mask',
    }

    # Get predictions
    pred = model(img)
    pred = non_max_suppression(pred, conf_thres, iou_thres, None, False, max_det=max_det)
    re_boxes = list()
    re_confidence = list()
    re_classes = list()
    re_mask_id = list()
    for i, det in enumerate(pred):  # per image
        det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], im.shape).round()
        for *xyxy, conf, cls in reversed(det):
            xyxy_list = torch.tensor(xyxy).view(1, 4).view(-1).tolist()
            conf_list = conf.tolist()
            label = names[int(cls)]
            re_boxes.append([int(xyxy_list[0]), int(xyxy_list[1]), int(xyxy_list[2]- xyxy_list[0]), int(xyxy_list[3]- xyxy_list[1])])
            re_confidence.append(conf_list)
            re_classes.append('face')
            re_mask_id.append(int(cls))
    print("Bounding Boxes:\n", re_boxes)
    print("Confidences:\n", re_confidence)
    print("Classes:\n", re_classes)
    print("Mask IDs:\n", re_mask_id)
    return re_boxes, re_confidence, re_classes, re_mask_id