计算机系统应用教程网站

网站首页 > 技术文章 正文

手把手教物体检测——YOLOV4 物体检测 pytorch

btikc 2024-10-12 11:19:27 技术文章 8 ℃ 0 评论

摘要

YOLOV4在coco上面达到了43.5%AP ,在Tesla V100 上达到了65FPS。相比今年的其它模型,得分不算高,但是它不是通过提高输入图像的分辨率来提高得分的,而是改进网络结构。创新点主要有以下几个方面:

(1)输入端:这里指的创新主要是训练时对输入端的改进,主要包括Mosaic数据增强、cmBN、SAT自对抗训练。

(2)BackBone主干网络:将各种新的方式结合起来,包括:CSPDarknet53、Mish激活函数、Dropblock

(3)Neck:目标检测网络在BackBone和最后的输出层之间往往会插入一些层,比如Yolov4中的SPP模块、FPN+PAN结构

(4)Prediction:输出层的锚框机制和Yolov3相同,主要改进的是训练时的损失函数CIOU_Loss,以及预测框筛选的nms变为DIOU_nms

训练

本地环境:

Ubuntu20.04

pytorch1.5.1

CUDA10.1

Python3.7

下载代码

本文使用的代码是:

训练和测试推理代码都已经完成。

下载权重文件:

yolov4.conv.137.pth

链接一: 提取码:kcel

链接二:


将下载的权重文件放到data文件夹下面

制作数据集

将Labelme数据集复制到pytorch-YOLOv4-master文件夹下面,如图:

然后用pycharm新建labelme2txt.py文件。写入生成训练集和验证集的代码。

labelme2txt.py代码:

from os import getcwdfrom sklearn.model_selection import train_test_splitimport jsonimport globwd = getcwd()"labelme标注的json 数据集转为pytorch版yolov4的训练集"classes = ["aircraft","oiltank"]image_ids = glob.glob(r"LabelmeData/*jpg")print(image_ids)train_list_file = open('data/train.txt', 'w')val_list_file = open('data/val.txt', 'w')def convert_annotation(image_id, list_file): jsonfile=open('%s.json' % (image_id)) in_file = json.load(jsonfile) for i in range(0,len(in_file["shapes"])): object=in_file["shapes"][i] cls=object["label"] points=object["points"] xmin=int(points[0][0]) ymin=int(points[0][1]) xmax=int(points[1][0]) ymax=int(points[1][1]) if cls not in classes: print("cls not in classes") continue cls_id = classes.index(cls) b = (xmin, ymin, xmax, ymax) list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id)) jsonfile.close()def ChangeData2TXT(image_List,dataFile): for image_id in image_List: dataFile.write('%s' % (image_id.split('\\')[-1])) convert_annotation(image_id.split('.')[0], dataFile) dataFile.write('\n') dataFile.close()trainval_files, test_files = train_test_split(image_ids, test_size=0.2, random_state=55)ChangeData2TXT(trainval_files,train_list_file)ChangeData2TXT(test_files,val_list_file)

安装代码需要的包。

参照requirements.txt安装本机没有的包,版本不一定要保持一致,只要后期不报错就没有事。

修改类别。

将coco.names和voc.names里面的类别修改为自己数据集的类别(默认是coco.names,都改了肯定没有错。),顺序和labelme2txt.py中的classes顺序保持一致。

修改配置文件cfg.py。

Cfg.use_darknet_cfg = False

Cfg.batch = 2(根据自己的显卡修改,我的显卡是8G的最多可以训练里2个batch)。

Cfg.subdivisions = 1

修改models.py

将51行和53行的inplace=True改为inplace=False。如果不修改,训练的时候会报个错误。

修改train.py文件。

找到526行,这个方法的参数是对cfg.py里面参数的更新。

主要修改的参数如下:

parser.add_argument('-g', '--gpu', metavar='G', type=str, default='0', help='GPU', dest='gpu')#设置GPU使用的GPU

parser.add_argument('-dir', '--data-dir', type=str, default="LabelmeData", help='dataset dir', dest='dataset_dir')#图片所在的文件夹。

parser.add_argument('-pretrained', type=str, default="data/yolov4.conv.137.pth", help='pretrained yolov4.conv.137')#设置预训练权重文件的路径。

parser.add_argument('-classes', type=int, default=80, help='dataset classes')#物体类别数。

parser.add_argument('-train_label_path', dest='train_label', type=str, default='data/train.txt', help="train label path")#训练集存放的路径。

注释415行到440行的代码,这段代码在验证的时候一直报错,我找不到原因。后续找到原因再更新。

将以上的内容修改完成后就可以点击run开始训练了。

测试

测试主要修改models.py的代码。将下面的代码从449行替换。


if __name__ == "__main__":    import sys    import cv2    namesfile = None    n_classes=2    weightfile="checkpoints/Yolov4_epoch151.pth"     imgfile="data/aircraft_4.jpg"#待测试的图片    width=608    height=608    model = Yolov4(yolov4conv137weight=None, n_classes=n_classes, inference=True)    pretrained_dict = torch.load(weightfile, map_location=torch.device('cpu'))#如果使用GPU则改为:#pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda'))

model.load_state_dict(pretrained_dict) use_cuda = True if use_cuda: model.cuda() img = cv2.imread(imgfile) # Inference input size is 416*416 does not mean training size is the same # Training size could be 608*608 or even other sizes # Optional inference sizes: # Hight in {320, 416, 512, 608, ... 320 + 96 * n} # Width in {320, 416, 512, 608, ... 320 + 96 * m} sized = cv2.resize(img, (width, height)) sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB) from tool.utils import load_class_names, plot_boxes_cv2 from tool.torch_utils import do_detect for i in range(2): # This 'for' loop is for speed check # Because the first iteration is usually longer boxes = do_detect(model, sized, 0.4, 0.6, use_cuda) if namesfile == None: if n_classes == 2: namesfile = 'data/voc.names' elif n_classes == 80: namesfile = 'data/coco.names' else: print("please give namefile") class_names = load_class_names(namesfile) resultImg=plot_boxes_cv2(img, boxes[0], 'predictions.jpg', class_names) cv2.imshow("image",resultImg) cv2.waitKey(0);

测试结果:

参考文章:

深入浅出Yolo系列之Yolov3&Yolov4核心基础知识完整讲解

https://blog.csdn.net/nan355655600/article/details/106246625?utm_medium=distribute.pc_relevant.none-task-blog-BlogCommendFromMachineLearnPai2-2.nonecase&depth_1-utm_source=distribute.pc_relevant.none-task-blog-BlogCommendFromMachineLearnPai2-2.nonecase

本文暂时没有评论,来添加一个吧(●'◡'●)

欢迎 发表评论:

最近发表
标签列表