旋转目标检测复现-yolov5-obb

复现源码：
https://github.com/hukaixuan19970627/yolov5_obb
亲测可行
安装流程：
按照https://github.com/hukaixuan19970627/yolov5_obb/blob/master/docs/install.md
确保安装过程不报错，否则影响后续训练

安装成功即可准备数据集


hf_txt存放划分好的训练集、测试集、验证集，里面内容为数据图像文件名，
images存放要训练的图像
labelTxt存放将xml转换后的txt标签文件
hf.py数据集划分;

# -*- coding: utf-8 -*-import osimport randomtrainval_percent = 0.9train_percent = 0.9xmlfilepath = 'xml'txtsavepath = 'images'total_xml = os.listdir(xmlfilepath)num = len(total_xml)list = range(num)tv = int(num * trainval_percent)tr = int(tv * train_percent)trainval = random.sample(list, tv)train = random.sample(trainval, tr)ftrainval = open('hf_txt/trainval.txt', 'w')ftest = open('hf_txt/test.txt', 'w')ftrain = open('hf_txt/train.txt', 'w')fval = open('hf_txt/val.txt', 'w')for i in list:    name = total_xml[i][:-4] + '\n'    if i in trainval:        ftrainval.write(name)        if i in train:            ftrain.write(name)        else:            fval.write(name)    else:        ftest.write(name)ftrainval.close()ftrain.close()fval.close()ftest.close()

xml转txt；

# 文件名称   ：roxml_to_dota.py# 功能描述   ：把rolabelimg标注的xml文件转换成dota能识别的xml文件，#             再转换成dota格式的txt文件#            把旋转框 cx,cy,w,h,angle，转换成四点坐标x1,y1,x2,y2,x3,y3,x4,y4import osimport xml.etree.ElementTree as ETimport mathdef edit_xml(xml_file,dotaxml_file):    """    修改xml文件    :param xml_file:xml文件的路径    :return:    """    tree = ET.parse(xml_file)    objs = tree.findall('object')    for ix, obj in enumerate(objs):        x0 = ET.Element("x0")  # 创建节点        y0 = ET.Element("y0")        x1 = ET.Element("x1")        y1 = ET.Element("y1")        x2 = ET.Element("x2")        y2 = ET.Element("y2")        x3 = ET.Element("x3")        y3 = ET.Element("y3")        # obj_type = obj.find('bndbox')        # type = obj_type.text        # print(xml_file)        if (obj.find('robndbox') == None):            obj_bnd = obj.find('bndbox')            obj_xmin = obj_bnd.find('xmin')            obj_ymin = obj_bnd.find('ymin')            obj_xmax = obj_bnd.find('xmax')            obj_ymax = obj_bnd.find('ymax')            xmin = float(obj_xmin.text)            ymin = float(obj_ymin.text)            xmax = float(obj_xmax.text)            ymax = float(obj_ymax.text)            obj_bnd.remove(obj_xmin)  # 删除节点            obj_bnd.remove(obj_ymin)            obj_bnd.remove(obj_xmax)            obj_bnd.remove(obj_ymax)            x0.text = str(xmin)            y0.text = str(ymax)            x1.text = str(xmax)            y1.text = str(ymax)            x2.text = str(xmax)            y2.text = str(ymin)            x3.text = str(xmin)            y3.text = str(ymin)        else:            obj_bnd = obj.find('robndbox')            obj_bnd.tag = 'bndbox'  # 修改节点名            obj_cx = obj_bnd.find('cx')            obj_cy = obj_bnd.find('cy')            obj_w = obj_bnd.find('w')            obj_h = obj_bnd.find('h')            obj_angle = obj_bnd.find('angle')            cx = float(obj_cx.text)            cy = float(obj_cy.text)            w = float(obj_w.text)            h = float(obj_h.text)            angle = float(obj_angle.text)            obj_bnd.remove(obj_cx)  # 删除节点            obj_bnd.remove(obj_cy)            obj_bnd.remove(obj_w)            obj_bnd.remove(obj_h)            obj_bnd.remove(obj_angle)            x0.text, y0.text = rotatePoint(cx, cy, cx - w / 2, cy - h / 2, -angle)            x1.text, y1.text = rotatePoint(cx, cy, cx + w / 2, cy - h / 2, -angle)            x2.text, y2.text = rotatePoint(cx, cy, cx + w / 2, cy + h / 2, -angle)            x3.text, y3.text = rotatePoint(cx, cy, cx - w / 2, cy + h / 2, -angle)        # obj.remove(obj_type)  # 删除节点        obj_bnd.append(x0)  # 新增节点        obj_bnd.append(y0)        obj_bnd.append(x1)        obj_bnd.append(y1)        obj_bnd.append(x2)        obj_bnd.append(y2)        obj_bnd.append(x3)        obj_bnd.append(y3)        tree.write(dotaxml_file, method='xml', encoding='utf-8')  # 更新xml文件# 转换成四点坐标def rotatePoint(xc, yc, xp, yp, theta):    xoff = xp - xc;    yoff = yp - yc;    cosTheta = math.cos(theta)    sinTheta = math.sin(theta)    pResx = cosTheta * xoff + sinTheta * yoff    pResy = - sinTheta * xoff + cosTheta * yoff    return str(int(xc + pResx)), str(int(yc + pResy))def totxt(xml_path,out_path):        # 想要生成的txt文件保存的路径，这里可以自己修改    files = os.listdir(xml_path)    for file in files:        tree = ET.parse(xml_path + os.sep + file)        root = tree.getroot()        name = file.strip('.xml')        output = out_path + name + '.txt'        file = open(output, 'w')        objs = tree.findall('object')        for obj in objs:            cls = obj.find('name').text            box = obj.find('bndbox')            x0 = int(float(box.find('x0').text))            y0 = int(float(box.find('y0').text))            x1 = int(float(box.find('x1').text))            y1 = int(float(box.find('y1').text))            x2 = int(float(box.find('x2').text))            y2 = int(float(box.find('y2').text))            x3 = int(float(box.find('x3').text))            y3 = int(float(box.find('y3').text))            file.write("{} {} {} {} {} {} {} {} {} 0\n".fORMat(x0, y0, x1, y1, x2, y2, x3, y3, cls))        file.close()        print(output)if __name__ == '__main__':    # -----**** 第一步：把xml文件统一转换成旋转框的xml文件 ****-----    roxml_path = "/root/autodl-tmp/yolov5_obb/dataset/dataset_demo/xml/"  # 目录下保存的是需要转换的xml文件    dotaxml_path = '/root/autodl-tmp/yolov5_obb/dataset/dataset_demo/1xml/'    out_path = '/root/autodl-tmp/yolov5_obb/dataset/dataset_demo/labelTxt/'    filelist = os.listdir(roxml_path)    for file in filelist:        edit_xml(os.path.join(roxml_path, file), os.path.join(dotaxml_path, file))    # -----**** 第二步：把旋转框xml文件转换成txt格式 ****-----    totxt(dotaxml_path, out_path)

voc_label.py划分训练集，验证集，测试集路径：

# -*- coding: utf-8 -*-import xml.etree.ElementTree as ETimport osfrom os import getcwdsets = ['train', 'val', 'test']classes = ["large_car","small_car"]abs_path = os.getcwd()wd = getcwd()for image_set in sets:    if not os.path.exists('labelTxt/'):        os.makedirs('labelTxt/')    image_ids = open('hf_txt/%s.txt' % (image_set)).read().strip().split()    list_file = open('%s.txt' % (image_set), 'w')    for image_id in image_ids:        list_file.write(abs_path + '/images/%s.jpg\n' % (image_id))    list_file.close()

修改相应数据集访问路劲

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]#path: ./dataset # dataset root dirtrain: dataset/dataset_demo/train.txt #images   # train images (relative to 'path') val: dataset/dataset_demo/val.txt #images  # val images (relative to 'path') #test: dataset_demo/images  #images # test images (optional)# Classesnc: 2  # number of classesnames: ['large_car','small_car']  # class names# Download script/URL (optional)# download: https://ultralytics.com/assets/coco128.zip

最后修改训练文件train.py,修改成对应的文件路径即可

上述都没问题即可训练

python train.py

扩展：部署yolov5-obb：
https://blog.csdn.net/qq_41043389/article/details/127777272

来源地址：https://blog.csdn.net/wxy2020915/article/details/128542603