pytorch下yolov8模型实现目标检测（全网最简洁快速，一眼懂）

最近研究了一个目标检测项目，需要对水稻进行检测，yolov8的模型官网都封装好了，官网地址：Home - Ultralytics YOLO Docs，本文给急需目标检测的同学参考，非常急的那种，废话少说，下面是所有代码：

import os
import torch
from ultralytics import YOLO

os.environ['KMP_DUPLICATE_LIB_OK']='True'
def main():

    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(device)

    model = YOLO('yolov8n.pt').to(device)

    model.train(data='safehat.yaml', epochs=100, device=device, workers=0,batch=2)

    model.val(device=device)

if __name__ == '__main__':
    main()

还需要一个yolov8n.pt的模型文件，这个官网可以下载我就不给了，下载后把官网文件中的yolov8n.pt拿出来和代码放到同一文件夹下面就可以了，记得安装ultralytics库，该代码搭配一个配置文件safehat.yaml（修改路径就可以使用，重要的只有那个路径）：

train: C:\Users\Desktop\ricedata\train
val: C:\Users\Desktop\ricedata\valid
test: C:\Users\Desktop\ricedata\test

nc: 3
#name: [Hardhat,Mask,NO-Hardhat,NO-Mask,NO-Safety Vest,Person,Safety Cone,Safety Vest,machinery,vehicle]
classes: [rice,animal,other]

最后是数据集的路径格式，他会自动获取符合规范路径下的数据，所以记得写好路径（.cache文件运行后才有）：

数据集都是我自己做的，标签名和图片名需要一致，推荐一个制作数据集的网站Newtab (makesense.ai)，运行后会自动在目录下生产文件夹，包含训练完的模型和训练过程的数据：

使用模型：

训练完成后会生成runs文件夹，包含两个训练文件夹，在其中一个里面找到best.pt的模型，这就是训练最好的目标检测模型（似乎还有个训练最差的模型）

import matplotlib.pyplot as plt
from ultralytics import YOLO

#加载已经训练好的模型
model = YOLO('runs/detect/train/weights/best.pt')

#选择要测试的图像路径
image_path = 'ricedata/test/images/0292.jpg'

results = model(image_path)
print(type(results))
print(results[0].orig_img)

plt.imshow(results[0].orig_img)
plt.show()
if isinstance(results, list):
    for result in results:
        # 显示检测结果
        result.show()
else:
    # 显示检测结果
   results.show()

下一篇文章我会对模型进行onnx打包，后续部署到安卓app上面