本文主要是介绍yolov8训练野火烟雾识别检测模型,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!
1.数据集下载
数据集下载链接:https://hyper.ai/datasets/33096
2. 数据集格式转换
需要将json中的标注信息转换为yolo格式的标注文件数据
import json
import os
import shutil
import cv2
import matplotlib.pyplot as plttarget = "./data/val"
def convert(size, box):dw = size[1]dh = size[0]# box x1 y1 x2 y2x = (box[0] + box[2]) / 2.0y = (box[1] + box[3]) / 2.0w = box[2] - box[0]h = box[3] - box[1]x = x / dww = w / dwy = y / dhh = h / dhif w >= 1:w = 0.99if h >= 1:h = 0.99return (x, y, w, h)
# 将标注数据转换为yolo格式
with open(target+"/_annotations.coco.json") as f:anno = json.load(f)images = {}labels = {}for img in anno['images']:images[img["id"]] = img["file_name"]for an in anno['annotations']:labels[an["image_id"]] = anprint(anno)img_dir = target+"/images/"anno_dir = target+"/labels/"if (not os.path.exists(img_dir)):os.mkdir(img_dir)os.mkdir(anno_dir)for i in images:# 将图片复制到images文件夹shutil.copyfile(target+"/"+ images[i], img_dir+"/"+ images[i])img = cv2.imread(img_dir + "/" + images[i])# 生成标注文件label = labels[i]filename,_ = os.path.splitext(images[i])with open(anno_dir+"/"+ filename+ ".txt","w") as f:box = label["bbox"]# img = cv2.rectangle(img,(box[0],box[1]),(box[0]+box[2],box[1] + box[3]),(50,50,50),2)# plt.imshow(img,)# plt.show()box = convert(img.shape, (box[0],box[1],box[0]+box[2],box[1] + box[3]))f.write(str(label["category_id"])+" " + " ".join([str(a) for a in box]))
将test、train和val都 转换一下
3. 模型训练
数据配置文件
# 数据集所在路径
path: C:\Users\lhq\Desktop\Wildfire-Smoke\datatrain: "./train/"
val: "./val/"
test: "./test/"nc: 2names:0: 烟雾1: 烟雾
训练代码
from ultralytics import YOLO
from ultralytics.utils import DEFAULT_CFG
from datetime import datetimecurrent_time = datetime.now()
time_str = current_time.strftime("%Y-%m-%d_%H-%M-%S")
# 训练结果保存路径
DEFAULT_CFG.save_dir = f"./models/{time_str}"if __name__ == "__main__":model = YOLO("yolov8n.pt") # Train the modelresults = model.train(data="smoke.yaml", epochs=100, imgsz=640, device=0, save=True)
4. 模型测试
预测代码:
from ultralytics import YOLO# Load a model
model = YOLO('best.pt')
# Run batched inference on a list of images
model.predict("./demo/", imgsz=640, save=True, device=0,plots=True)
这篇关于yolov8训练野火烟雾识别检测模型的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!