YOLO-物体检测

#encoding:utf-8
from ultralytics import YOLO
from PIL import Image, ImageDraw,ImageFont

model  = YOLO("yolov8n.pt")
image1_path = "img/guangzhou.jpg"
result = model.predict(image1_path)

img = Image.open(image1_path)
draw = ImageDraw.Draw(img)
font_size = 16
font = ImageFont.truetype("/usr/share/fonts/truetype/wqy/wqy-microhei.ttc",size=font_size)

classify_dict = {}
for r in result:
    names = r.names
    boxes = r.boxes
    probs = boxes.conf
    labels = boxes.cls
    label_num = labels.shape[0]
    for i in range(label_num):
        box = boxes[i].xyxy[0].cpu().numpy()
        name = names[int(labels[i])]
        probability = round(float(probs[i]),2)
        if name not in classify_dict:
            num = 1
            classify_dict[name] = num
        else:
            num = classify_dict[name] + 1
            classify_dict[name] += 1
        if probability >= 0.3:
            draw.rectangle(box,outline="red",width=1)
            half_width = int(box[0]+box[2])//2
            half_height = int(box[1]+box[3])//2
            text_position = (half_width-font_size*2,half_height)
            draw.rectangle((box[0],half_height,box[2],half_height+font_size+4),fill="white")
            text = "{}{}: {}".format(name, num,probability)
            draw.text(text_position,text,font=font,fill="black")
img.save("out.jpg")

最终效果如下:

模型下载可以参考
模型下载:https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.pt
参考链接:

https://docs.ultralytics.com/modes/predict/