Implementing YOLOv8

使用Ultralytics提供的方法。

---

# 序

不同于以往将模型写在主函数里，此处依赖\(\text{Ultralytics}\)提供的一系列方法使用\(\text{YOLOv8}\)模型。

# 安装依赖

官方推荐使用 pip 进行该库的安装。

pip install ultralytics
#使用代理
pip install --proxy http://127.0.0.1:2081 ultralytics

# 使用预训练模型

\(\text{Ultralytics}\)极大简化了使用预训练模型直接进行推理的流程。

实例化预训练好的\(\text{yolov8n}\)

model = YOLO("yolov8n.pt")

定义图像文件夹路径等基础参数

img_dir = os.path.join('/workspace/data/2024/')
trans = v2.Compose([v2.ToTensor()])
imgs_list = ['test.jpg', 'test2.jpg']
imgs = []

加载列表内指定的图像

for file_name in imgs_list:
    imgs.append(torchvision.io.read_image(f'{img_dir}{file_name}', mode=torchvision.io.image.ImageReadMode.RGB))

使用\(\text{YOLOv8}\)推理

def img_name2img_location(imgs_list:list, img_dir:str):
    imgs_location_list = [f'{img_dir}{img_item}' for img_item in imgs_list]
    return imgs_location_list

yolo_rslts = model(img_name2img_location(imgs_list=imgs_list, img_dir=img_dir), conf=0.335)

根据推理结果输出图片

def yolo_xywh2patch_xywh(xywh:list):
    # yolo的xy是bbox的中心点xy
    xywh_mtpltlb = []
    for xywh_item in xywh:
        xywh_item_list = []
        xywh_item_list.append(xywh_item[0] - (xywh_item[2] + 100) / 2)
        xywh_item_list.append(xywh_item[1] - (xywh_item[3] + 100) / 2)
        xywh_item_list.append(xywh_item[2] + 100)
        xywh_item_list.append(xywh_item[3] + 100)
        xywh_mtpltlb.append(xywh_item_list)
    return xywh_mtpltlb

for idx, result in enumerate(yolo_rslts):
    # 根据每张图像比例实例化新的subplot，这样才能配合隐藏坐标轴等参数仅输出图片本身
    w, h = matplotlib.figure.figaspect(imgs[idx].shape[1] / imgs[idx].shape[2])
    fig, ax = pyplot.subplots(figsize=(w, h))
    boxes = result.boxes  # Boxes object for bounding box outputs
    xywh_mtpltlb = yolo_xywh2patch_xywh(boxes.xywh.to('cpu'))
    print(boxes)
    masks = result.masks  # Masks object for segmentation masks outputs
    keypoints = result.keypoints  # Keypoints object for pose outputs
    probs = result.probs  # Probs object for classification outputs
    obb = result.obb  # Oriented boxes object for OBB outputs
    # result.show()  # display to screen
    # result.save(filename=f'result_{idx}.jpg')  # save to disk
    ax.imshow(imgs[idx].permute(1, 2, 0))
    rect_list = []
    for i in range(len(boxes.xywh)):
        rect_list.append(patches.Rectangle(xy=(xywh_mtpltlb[i][0], xywh_mtpltlb[i][1]),
                                           width=xywh_mtpltlb[i][2], height=xywh_mtpltlb[i][3],
                                           color='red', fc='none', ec='blue', lw=.3))
        ax.text(xywh_mtpltlb[i][0], xywh_mtpltlb[i][1], f'Person: {boxes.conf.to("cpu")[i]:.2f}', bbox=dict(fc='white', alpha=.6, ec='none', boxstyle='Round, pad=0.1'), size=5)
    for rect in rect_list:
        ax.add_patch(rect)
    ax.axis('off')
    pyplot.tight_layout(pad=0)
    pyplot.savefig(f'testsave_{idx}.jpg', dpi=600)

注：\(\text{pyplot.tight_layout(pad=0)}\)和\(\text{ax.axis(‘off’)}\)需要配合与图像一致的\(\text{Figure}\)比例的时候才会不输出白色边框背景。