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目录：存放测试集图片的标签

训练代码

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'))

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

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")