调用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