|
- import numpy as np
- import torch
-
- from utils.datasets import letterbox
- from utils.general import non_max_suppression, overlap_box_suppression, scale_coords
-
-
- def yolo_process(img0, model, device, args, half):
- # Padded resize
- stride = int(model.stride.max()) # model stride
-
- img, ratio, (dw, dh) = letterbox(img0, args.yoloimg_size, stride=stride)
-
- # Convert
- img = img[:, :, ::-1].transpose(2, 0, 1) # BGR to RGB, to 3x416x416
- img = np.ascontiguousarray(img)
- img = torch.from_numpy(img).to(device)
- img = img.half() if half else img.float() # uint8 to fp16/32
- img /= 255.0 # 0 - 255 to 0.0 - 1.0
- if img.ndimension() == 3:
- img = img.unsqueeze(0)
-
- pred = model(img, augment=args.augment)[0]
- # Apply NMS
- pred = non_max_suppression(pred, args.conf_thres, args.iou_thres, classes=args.classes, agnostic=args.agnostic_nms)
- pred = overlap_box_suppression(pred, args.ovlap_thres)
-
- ###一次检测一张图片
- det = pred[0]
- if len(det):
- # Rescale boxes from img_size to im0 size
- det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()
-
- return det
|
