yolov5格式的香烟数据集

https://download.csdn.net/download/qq_42864343/88110620?spm=1001.2014.3001.5503

创建yolo-nas的运行环境

进入Pycharm的terminal，输入如下命令

conda create -n yolonas python=3.8pip install super-gradients

使用自定义数据训练Yolo-nas

准备数据

在YOLO-NAS根目录下创建mydata文件夹（名字可以自定义），目录结构如下：

将自己数据集里用labelImg标注好的xml文件放到xml目录
图片放到images目录

划分数据集

把划分数据集代码 split_train_val.py放到yolo-nas目录下:

# coding:utf-8import os
import random
import argparse# 通过argparse模块创建一个参数解析器。该参数解析器可以接收用户输入的命令行参数，用于指定xml文件的路径和输出txt文件的路径。
parser = argparse.ArgumentParser()
# 指定xml文件的路径
parser.add_argument('--xml_path', default='mydata/xml', type=str, help='input xml label path')
# 设置输出txt文件的路径
parser.add_argument('--txt_path', default='mydata/dataSet', type=str, help='output txt label path')
opt = parser.parse_args()
# 训练集与验证集 占全体数据的比例
trainval_percent = 1.0
# 训练集 占训练集与验证集总体 的比例
train_percent = 0.9
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
# 获取到xml文件的数量
total_xml = os.listdir(xmlfilepath)
# 判断txtsavepath是否存在，若不存在，则创建该路径。
if not os.path.exists(txtsavepath):os.makedirs(txtsavepath)# 统计xml文件的个数，即Image标签的个数
num = len(total_xml)
list_index = range(num)
# tv (训练集和测试集的个数)　＝　数据总数　＊　训练集和数据集占全体数据的比例
tv = int(num * trainval_percent)
# 训练集的个数
tr = int(tv * train_percent)
# 　按数量随机得到取训练集和测试集的索引
trainval = random.sample(list_index, tv)
# 　打乱训练集　
train = random.sample(trainval, tr)
# 　创建存放所有图片数据路径的文件
file_trainval = open(txtsavepath + '/trainval.txt', 'w')
# 　创建存放所有测试图片数据的路径的文件
file_test = open(txtsavepath + '/test.txt', 'w')
# 创建存放所有训练图片数据的路径的文件
file_train = open(txtsavepath + '/train.txt', 'w')
# 创建存放所有测试图片数据的路径的文件
file_val = open(txtsavepath + '/val.txt', 'w')# 遍历list_index列表，将文件名按照划分规则写入相应的txt文件中
for i in list_index:name = total_xml[i][:-4] + '\n'if i in trainval:file_trainval.write(name)if i in train:file_train.write(name)else:file_val.write(name)else:file_test.write(name)file_trainval.close()
file_train.close()
file_val.close()
file_test.close()

运行代码:
dataSet中出现四个文件，里面是图片的名字

根据xml标注文件制作适合yolo的标签

即将每个xml标注提取bbox信息为txt格式，每个图像对应一个txt文件，文件每一行为一个目标的信息，包括class, x_center, y_center, width, height。
创建make_labes.py，复制如下代码运行：

# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import os
from os import getcwdsets = ['train', 'val', 'test']
classes = ['smoke']   # 改成自己的类别
abs_path = os.getcwd()
print(abs_path)def convert(size, box):dw = 1. / (size[0])dh = 1. / (size[1])x = (box[0] + box[1]) / 2.0 - 1y = (box[2] + box[3]) / 2.0 - 1w = box[1] - box[0]h = box[3] - box[2]x = x * dww = w * dwy = y * dhh = h * dhreturn x, y, w, hdef convert_annotation(image_id):in_file = open('mydata/xml/%s.xml' % (image_id), encoding='UTF-8')out_file = open('mydata/label/%s.txt' % (image_id), 'w')tree = ET.parse(in_file)root = tree.getroot()size = root.find('size')w = int(size.find('width').text)h = int(size.find('height').text)for obj in root.iter('object'):difficult = obj.find('difficult').textcls = obj.find('name').textif cls not in classes or int(difficult) == 1:continuecls_id = classes.index(cls)xmlbox = obj.find('bndbox')b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),float(xmlbox.find('ymax').text))b1, b2, b3, b4 = b# 标注越界修正if b2 > w:b2 = wif b4 > h:b4 = hb = (b1, b2, b3, b4)bb = convert((w, h), b)out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')wd = getcwd()
for image_set in sets:if not os.path.exists('mydata/label/'):os.makedirs('mydata/label/')image_ids = open('mydata/dataSet/%s.txt' % (image_set)).read().strip().split()list_file = open('mydata/%s.txt' % (image_set), 'w')for image_id in image_ids:list_file.write(abs_path + '/mydata/images/%s.jpg\n' % (image_id))convert_annotation(image_id)list_file.close()

运行完成：

label目录下出现了图片对应的标记位置(好像是标记框的左上角和由上角的坐标)与类别

mydata目录下，出现了训练集train.txt，测试集test.txt,里面是对应的图片路径

将划分好的数据集转成适合yolo-nas要求的数据集

创建data目录

error目录: 存放格式有问题的图片，格式有问题的图片会中断训练
images/train目录：存放训练集图片
images/val目录：存放测试集图片
labels/train目录：存放训练集图片的标签
labels/val目录：存放测试集图片的标签

根据 trian.txt和val.txt中的图片路径移动图片到train和val文件中

项目目录下创建 mv_img.py，运行代码移动图片

import shutil# 读取txt文件根据里面的图片路径移动图片到新目录
def move_img(img_txt_path,new_img_path):with open(img_txt_path, 'r') as file:image_paths = file.readlines()# image_new_paths = []# for path in image_paths:#     path = path.strip()#     path = 'mydata/label/'+path+'.txt'#     image_new_paths.append(path)# 去除路径中的换行符image_paths = [path.strip() for path in image_paths]# 复制图片到新的文件夹# for path in image_new_paths:for path in image_paths:shutil.copy(path, new_img_path)
img_txt_paths =['mydata/val.txt','mydata/train.txt']
new_img_paths =['data/images/val','data/images/train']for img_txt_path,new_img_path in zip(img_txt_paths,new_img_paths):print(f'img_txt_path,new_img_path = {img_txt_path,new_img_path}')move_img(img_txt_path,new_img_path)
# move_img('mydata/train.txt','data/images/train')

项目目录下创建 mv_labels.py，运行代码移动标签

import shutil
# 根据图片名，读取标签的txt文件，并移动到data中label文件夹中
def move_label(img_txt_path,new_label_path):with open(img_txt_path, 'r') as file:image_paths = file.readlines()image_new_paths = []for path in image_paths:path = path.strip()path = 'mydata/label/'+path+'.txt'image_new_paths.append(path)print(f'image_new_paths = {image_new_paths}')# 去除路径中的换行符# image_paths = [path.strip() for path in image_paths]# 复制图片到新的文件夹# for path in image_new_paths:for path in image_new_paths:print(f'path={path}')print(f'new_label_path={new_label_path}')shutil.copy(path, new_label_path)'''移动标签'''
img_name_paths =['mydata/dataSet/val.txt','mydata/dataSet/train.txt']
new_label_paths =['data/labels/val','data/labels/train']for img_name_path,new_label_path in zip(img_name_paths,new_label_paths):print(f'img_name_path,new_img_path = {img_name_path,new_label_path}')move_label(img_name_path,new_label_path)

去除训练集测试集中格式错误的代码

在项目的根目录下创建 检查错误jpg.py，复制如下代码运行，这样会去除掉data/images/val/路径下格式错误的图片(错误图片会使训练中的程序中断运行)

import ostrain_dir = 'data/images/val/'def progress(percent, width=50):'''进度打印功能'''if percent >= 100:percent = 100show_str = ('[%%-%ds]' % width) % (int(width * percent / 100) * "#")  # 字符串拼接的嵌套使用print('\r%s %d%% ' % (show_str, percent), end='')def is_valid_jpg(jpg_file):with open(jpg_file, 'rb') as f:f.seek(-2, 2)buf = f.read()f.close()return buf == b'\xff\xd9'  # 判定jpg是否包含结束字段data_size = len([lists for lists in os.listdir(train_dir) if os.path.isfile(os.path.join(train_dir, lists))])
recv_size = 0
incompleteFile = 0
print('file tall : %d' % data_size)for file in os.listdir(train_dir):if os.path.splitext(file)[1].lower() == '.jpg':ret = is_valid_jpg(train_dir + file)if ret == False:incompleteFile = incompleteFile + 1os.remove(train_dir + file)recv_per = int(100 * recv_size / data_size)progress(recv_per, width=30)recv_size = recv_size + 1progress(100, width=30)
print('\nincomplete file : %d' % incompleteFile)

然后将train_dir = 'data/images/val/'改为train_dir = ‘data/images/train/’,再运行一下代码，去除掉train文件下格式错误的图片

训练代码

import osimport requests
import torch
from PIL import Imagefrom super_gradients.training import Trainer, dataloaders, models
from super_gradients.training.dataloaders.dataloaders import (coco_detection_yolo_format_train, coco_detection_yolo_format_val
)
from super_gradients.training.losses import PPYoloELoss
from super_gradients.training.metrics import DetectionMetrics_050
from super_gradients.training.models.detection_models.pp_yolo_e import (PPYoloEPostPredictionCallback
)

class config:# trainer paramsCHECKPOINT_DIR = 'checkpoints'  # specify the path you want to save checkpoints toEXPERIMENT_NAME = 'cars-from-above'  # specify the experiment name# dataset paramsDATA_DIR = 'data'  # parent directory to where data livesTRAIN_IMAGES_DIR = 'images/train'  # child dir of DATA_DIR where train images areTRAIN_LABELS_DIR = 'labels/train'  # child dir of DATA_DIR where train labels areVAL_IMAGES_DIR = 'images/val'  # child dir of DATA_DIR where validation images areVAL_LABELS_DIR = 'labels/val'  # child dir of DATA_DIR where validation labels are# TEST_IMAGES_DIR = 'images/test'  # child dir of DATA_DIR where validation images are# TEST_LABELS_DIR = 'labels/test'  # child dir of DATA_DIR where validation labels areCLASSES = ['smoke']  # 指定类名NUM_CLASSES = len(CLASSES) # 获取类个数# dataloader params - you can add whatever PyTorch dataloader params you have# could be different across train, val, and testDATALOADER_PARAMS = {'batch_size': 16,'num_workers': 2}# model paramsMODEL_NAME = 'yolo_nas_l'  # 可以选择 yolo_nas_s, yolo_nas_m, yolo_nas_l。分别是 小型，中型，大型PRETRAINED_WEIGHTS = 'coco'  # only one option here: coco

trainer = Trainer(experiment_name=config.EXPERIMENT_NAME, ckpt_root_dir=config.CHECKPOINT_DIR)# 指定训练数据
train_data = coco_detection_yolo_format_train(dataset_params={'data_dir': config.DATA_DIR,'images_dir': config.TRAIN_IMAGES_DIR,'labels_dir': config.TRAIN_LABELS_DIR,'classes': config.CLASSES},dataloader_params=config.DATALOADER_PARAMS
)# 指定评估数据
val_data = coco_detection_yolo_format_val(dataset_params={'data_dir': config.DATA_DIR,'images_dir': config.VAL_IMAGES_DIR,'labels_dir': config.VAL_LABELS_DIR,'classes': config.CLASSES},dataloader_params=config.DATALOADER_PARAMS
)# test_data = coco_detection_yolo_format_val(
#     dataset_params={
#         'data_dir': config.DATA_DIR,
#         'images_dir': config.TEST_IMAGES_DIR,
#         'labels_dir': config.TEST_LABELS_DIR,
#         'classes': config.CLASSES
#     },
#     

dataloader_params=config.DATALOADER_PARAMS
# )
# train_data.dataset.plot()model = models.get(config.MODEL_NAME,num_classes=config.NUM_CLASSES,pretrained_weights=config.PRETRAINED_WEIGHTS)
train_params = {# ENABLING SILENT MODE"average_best_models":True,"warmup_mode": "linear_epoch_step","warmup_initial_lr": 1e-6,"lr_warmup_epochs": 3,"initial_lr": 5e-4,"lr_mode": "cosine","cosine_final_lr_ratio": 0.1,"optimizer": "Adam","optimizer_params": {"weight_decay": 0.0001},"zero_weight_decay_on_bias_and_bn": True,"ema": True,"ema_params": {"decay": 0.9, "decay_type": "threshold"},# ONLY TRAINING FOR 10 EPOCHS FOR THIS EXAMPLE NOTEBOOK"max_epochs": 200,"mixed_precision": True,"loss": PPYoloELoss(use_static_assigner=False,# NOTE: num_classes needs to be defined herenum_classes=config.NUM_CLASSES,reg_max=16),"valid_metrics_list": [DetectionMetrics_050(score_thres=0.1,top_k_predictions=300,# NOTE: num_classes needs to be defined herenum_cls=config.NUM_CLASSES,normalize_targets=True,post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.01,nms_top_k=1000,max_predictions=300,nms_threshold=0.7))],"metric_to_watch": 'mAP@0.50'
}trainer.train(model=model,training_params=train_params,train_loader=train_data,valid_loader=val_data)best_model = models.get(config.MODEL_NAME,num_classes=config.NUM_CLASSES,checkpoint_path=os.path.join(config.CHECKPOINT_DIR, config.EXPERIMENT_NAME, 'average_model.pth'))

训练时仍然代码检测不到的错误格式图片

根据第一行红字我们发现它的id是338。

根据id找到图片路径

import osdef print_file_name(directory, file_index):files = os.listdir(directory)# 确保文件索引在有效范围内if file_index >= 0 and file_index < len(files):file_name = files[file_index]print("第{}个文件的文件名为：{}".format(file_index + 1, file_name))else:print("无法找到第{}个文件".format(file_index + 1))# 指定目录路径
directory_path = "data/images/train"# 要打印的文件索引
file_index = 338# 调用函数打印指定文件的文件名
print_file_name(directory_path, file_index)

发现是：data/images/train/smoke_b002192.jpg

根据文件路径将图片移动到error文件夹

import os
import shutildef move_file_to_new_directory(file_name):# 源目录的路径train_dir = "data/images/train"# 目标目录的路径new_dir = "data/error"# 构造文件的完整路径file_path = os.path.join(train_dir, file_name)if os.path.exists(file_path):# 移动文件到new目录shutil.move(file_path, new_dir)print(f"成功将文件{file_name}移动到new目录")else:print(f"找不到文件{file_name}")# 测试 ，输入图片名，将图片移动
move_file_to_new_directory("smoke_b002192.jpg")

再次运行训练代码

成功开始训练。

连接网络摄像头用训练好的模型参数进行预测

import torch
from super_gradients.training import models
import cv2
import time
def get_video_capture(video, width=None, height=None, fps=None):"""获得视频读取对象--   7W   Pix--> width=320,height=240--   30W  Pix--> width=640,height=480720P,100W Pix--> width=1280,height=720960P,130W Pix--> width=1280,height=10241080P,200W Pix--> width=1920,height=1080:param video: video file or Camera ID:param width:   图像分辨率width:param height:  图像分辨率height:param fps:  设置视频播放帧率:return:"""video_cap = cv2.VideoCapture(video)# 如果指定了宽度，高度，fps，则按照制定的值来设置，此处并没有指定if width:video_cap.set(cv2.CAP_PROP_FRAME_WIDTH, width)if height:video_cap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)if fps:video_cap.set(cv2.CAP_PROP_FPS, fps)return video_cap# 此处连接网络摄像头进行测试
video_file = 'rtsp://账号:密码@ip/Streaming/Channels/1'
# video_file = 'data/output.mp4'
num_classes = 1
# best_pth = '/home/computer_vision/code/my_code/checkpoints/cars-from-above/ckpt_best.pth'
best_pth = 'checkpoints/cars-from-above/smoke_small_ckpt_best.pth'
device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
best_model = models.get("yolo_nas_s", num_classes=num_classes, checkpoint_path=best_pth).to(device)'''开始计时'''
start_time = time.time()
video_cap = get_video_capture(video_file)
while True:isSuccess, frame = video_cap.read()if not isSuccess:breakresult_image = best_model.predict(frame, conf=0.45, fuse_model=False)result_image = result_image._images_prediction_lst[0]result_image = result_image.draw()'''改动'''result_image = cv2.resize(result_image, (960, 540))'''end'''cv2.namedWindow('result', flags=cv2.WINDOW_NORMAL)cv2.imshow('result', result_image)kk = cv2.waitKey(1)if kk == ord('q'):break
video_cap.release()
'''时间结束'''
end_time = time.time()
run_time = end_time - start_time
print(run_time)

补充

对视频进行预测

import torch
from super_gradients.training import modelsdevice = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
model = models.get("yolo_nas_l", pretrained_weights="coco").to(device)
model.predict("data/output.mp4",conf=0.4).save("output/output_lianzhang.mp4")

对图片进行预测

import torch
from super_gradients.training import modelsdevice = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
model = models.get("yolo_nas_s", pretrained_weights="coco").to(device)
out = model.predict("camera01.png", conf=0.6)
out.show()
out.save("output")

预测data目录下的视频并保存预测结果

model.predict("data/output.mp4").save("output/output_lianzhang.mp4")