调用API方式1
首先同目录下需要复制YOLOv5代码的models文件、utils文件、export文件。 返回格式为列表形式，存储所有的目标位置，置信度等 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调用API方式1
首先同目录下需要复制YOLOv5代码的models文件、utils文件、export文件。 返回格式为列表形式，存储所有的目标位置，置信度等 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
