34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
|
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
|