pytorch_SSD训练评估自己的数据集

最新推荐文章于 2023-11-13 17:12:05 发布

Lighter_Go

最新推荐文章于 2023-11-13 17:12:05 发布

阅读量1.1k

点赞数 1

分类专栏：目标检测文章标签： python ssd 深度学习 pytorch

本文链接：https://blog.csdn.net/Lighter_Go/article/details/106917245

版权

本文介绍了如何在PyTorch中使用SSD模型训练自定义数据集，从Yolo格式转换为VOC格式，详细阐述了数据集处理、模型环境配置、代码修改以及解决训练过程中遇到的BUG问题。

摘要由CSDN通过智能技术生成

1. 模型环境

利用**ssd.pytorch**模型
运行环境显卡P5000 ubuntu16.04+CUDA9.2
pytorch1.0.1 , torchvision 0.2.2

2. 数据集处理

2.1 数据格式转换yolo_to_voc

由于开始制作的数据集是为了在yolov3模型下进行训练的yolo格式，所以我们需要进行转换。
原始数据： 640×640 yolo格式
目标数据： 300×300 voc格式

# -*- coding = utf-8 -*-

import copy
from lxml.etree import Element, SubElement, tostring, ElementTree
import cv2
import os

# 为了把yolo格式标签name_value转换成name
# 注意ssd模型中标签name英文字母需要小写
labels_dict = {
   "0":'a',"1":'b',"2":'c',"3":'d',"4":'e',"5":'f',"6":'g',"7":'h',"8":'i',
               "9":'j',...,"17":'z' }
cnt = 0
txt_path = ''  # yolo存放txt的文件目录
image_path = ''  # 存放图片的文件目录
path = ''  # 存放生成xml的文件目录

xml_head = '''<annotation>
    <folder>VOC2007</folder>
    <filename>{}</filename>
    <source>
        <database>Unknown</database>
    </source>   
    <size>
        <width>{}</width>
        <height>{}</height>
        <depth>{}</depth>
    </size>
    <segmented>0</segmented>
    '''
xml_obj = '''
    <object>        
        <name>{}</name>
        <pose>Rear</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>{}</xmin>
            <ymin>{}</ymin>
            <xmax>{}</xmax>
            <ymax>{}</ymax>
        </bndbox>
    </object>
    '''
xml_end = '''
</annotation>'''


for (root, dirname, files) in os.walk(image_path):  # 遍历图片文件夹
    for ft in files:
        ftxt = ft.replace('jpg', 'txt')  # ft是图片名字+扩展名，将jpg和txt替换
        fxml = ft.replace('jpg', 'xml')
        xml_path = path + fxml
        obj = ''

        img = cv2.imread(root + ft)
        img_h, img_w = img.shape[0], img.shape[1]
        head = xml_head.format(str(ft), str(img_w), str(img_h), 3)

        with open(txt_path + ftxt, 'r') as f:  # 读取对应txt文件内容
            for line in f.readlines():
                yolo_datas = line.strip().split(' '