YOLOv5最简推理分割代码

目录结构如下：

以加载torchscript模型文件为例，如下为Predict.py代码，utils.py中代码则从原始YOLOV5代码中拷贝，即Predict.py中用到的代码块，其中对部分用不到的代码进行了删减：

import cv2
import numpy as np
import torch

from utils.utils import letterbox, non_max_suppression, Annotator, scale_boxes, process_mask, colors
from pathlib import Path
# 路径设置
image_path = './data/1.jpg'
output_path = './runs/'
onnx_model_path = './weights/best.torchscript'

# 类别名及绘图颜色设置
names = ['crack']

# 硬件选择，优先选用GPU
device = torch.device('cuda:0')
# 1、加载模型,模型半精度
model = torch.jit.load(onnx_model_path, map_location=device)
model.half()  # 此处用cpu推理的话会报错，若用cpu，就把全局的half()删掉
# 2、图像预处理，图像半精度
img0 = cv2.imread(image_path)
img = letterbox(img0, new_shape=(416, 416), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32)[0]  # padded resize
img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
img = np.expand_dims(img, 0)          # 由于没有设置batchsize，用的单张图片，所以需要增大一个维度
img = np.ascontiguousarray(img)       # contiguous，换为内存连续存储的数组，使得运行速度更快

# 3、加载图像到GPU
img = torch.from_numpy(img).float().to(device)
img = img.half()  # fp16/32 img.float()
img /= 255  # 0 - 255 to 0.0 - 1.0

# 4、模型推理
pred, proto = model(img)
# 5、NMS
pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45, nm=32)
# 6、结果原图绘制
for i, det in enumerate(pred):
    annotator = Annotator(img0, line_width=3, example=str(names))
    if len(det):
        masks = process_mask(proto[i], det[:, 6:], det[:, :4], img.shape[2:], upsample=True)  # HWC
        det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], img0.shape).round()  # rescale boxes to im0 size

        annotator.masks(
            masks,
            colors=[colors(x, True) for x in det[:, 5]],
            im_gpu=img[0])
        for j, (*xyxy, conf, cls) in enumerate(reversed(det[:, :6])):
            c = int(cls)  # integer class
            label = f'{names[c]} {conf:.2f}'
            annotator.box_label(xyxy, label, color=colors(c, True))
    # Stream results
    im0 = annotator.result()
    save_image_name = Path(image_path).name
    cv2.imwrite(output_path + save_image_name, im0)

if __name__ == '__main__':
    print('hello')