YOLO检测结果渲染方法

代码示例

from ultralytics import YOLO
import cv2

image = cv2.imread("./test.jpg") 
def predict_image(image):
    model = YOLO("./model/best.pt")
    results = model.predict(source=image, conf=0.8, save=True)
    if resutls is not None and results.size>0:    
        show(results[0])

file_folder = "./test_images"
def predict_files(file_folder):
    model = YOLO("./model/best.pt")
    results = model.predict(source=file_folder, conf=0.8, save=True)
    for res in results:
        show(res)

def show(result):
    image = result.plot()
    for box, contour in zip(result.boxes.xyxy.cpu().numpy(), result.masks.xy):
        contour = np.array(contour, dtype=np.float32)
        box = np.array(box).astype(np.int32)
        M = cv2.moments(contour)
        if M["m00"]!=0:
            cx = int(M["m10"]/M["m00"])
            cy = int(M["m01"]/M["m00"])
        else:
            cx,cy=0,0
        x1,y1,x2,y2= box
        cv2.circle(image, (cx,cy), 5, (0,0,255), 2)
        cv2.putText(image, show_info, (cx, cy),cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,255,0),2)
    cv2.imwrite("results/result.jpg", image)

结果示例

关键代码说明

1. model.predict(source=xxx, conf=0.8, save=True) 参数source传入检测的目标，可以是一个opencv的cv图像，可以是一个图片文件的路径，可以是一个文件夹的路径 参数conf代表置信度阈值，即大于该值的检测结果会被认为是目标，否则会被过滤 参数save开关为True时，每次执行predict方法，会在run文件夹下生成一个predict子文件夹，存放检测结果文件

2. result.plot() 获取检测结果，会自动标注上目标的分类标签名称、置信度、目标包围盒位置

3. cv2.moments 计算质心 M = cv2.moments(contour) cx = int(M["m10"]/M["m00"]) cy = int(M["m01"]/M["m00"])