• 7.1 yolov5优化模型时,自动标注xml数据


    yolov5优化模型时,一般需要继续标注一些检测错误的图片,将其标为xml数据。以下是根据训练好的模型自动标注xml数据的python代码:

    注意:代码中包含了本人的yolov5的测试过程,测试过程可以自己根据yolov5的测试文件自行修改,只是测试返回的类格式为:

    [["water",[15,20,30,40]],["red",[12,13,14,15]]]

    二维数组表示测试的类为water和red,其中后面的数字表示类的坐标:[top,left,bottom,right],表示上、左、下、右4个坐标。

    1. import os
    2. import cv2
    3. from PIL import Image
    4. from yolo import YOLO
    5. #1.预测类,获得字符串
    6. class Predict():
    7. def a(self, img_path,save_path,img_name):
    8. image = Image.open(img_path)
    9. r_image, pred = yolo.detect_image(image, pred_class, img_name)
    10. if not os.path.exists(dir_save_path):
    11. os.makedirs(dir_save_path)
    12. r_image.save(save_path, quality=95, subsampling=0)
    13. return pred
    14. #2.写入xml文件
    15. def img_xml(img_path,xml_path,img_name,pred):
    16. if len(pred) != 0:
    17. #1.读取图片(xml需要写入图片的长宽高)
    18. img = cv2.imread(img_path)
    19. #2.写入xml文件
    20. #(1)写入文件头部
    21. files_path=img_path.split("\\")[-2]
    22. print("..:",files_path)
    23. xml_file = open((xml_path + img_name + '.xml'), 'w')
    24. xml_file.write('\n')
    25. xml_file.write(' ' +files_path+ '\n')
    26. xml_file.write(' ' + img_name + '.jpg' + '\n')
    27. xml_file.write(' ' + img_path +'\n')
    28. xml_file.write(' \n')
    29. xml_file.write(' Unknown\n')
    30. xml_file.write(' \n')
    31. #(2)写入图片的长宽高信息
    32. xml_file.write(' \n')
    33. xml_file.write(' '+str(img.shape[1])+'\n')
    34. xml_file.write(' ' + str(img.shape[0]) + '\n')
    35. xml_file.write(' ' + str(img.shape[2]) + '\n')
    36. xml_file.write(' \n')
    37. xml_file.write(' 0\n')
    38. #3.写入字符串信息:[["water",[15,20,30,40]],["red",[12,13,14,15]]]
    39. #if len(shuzu)!=0:
    40. for item in pred:
    41. xml_file.write(' \n')
    42. xml_file.write(' ' + str(item[0]) + '\n')
    43. xml_file.write(' Unspecified\n')
    44. xml_file.write(' 0\n')
    45. xml_file.write(' 0\n')
    46. xml_file.write(' \n')
    47. #写入字符串信息
    48. #[top, left, bottom, right]
    49. xml_file.write(' ' + str(item[1][1]) + '\n')
    50. xml_file.write(' ' + str(item[1][0]) + '\n')
    51. xml_file.write(' ' + str(item[1][3]) + '\n')
    52. xml_file.write(' ' + str(item[1][2]) + '\n')
    53. xml_file.write(' \n')
    54. xml_file.write('
    55. \n'
      )
    56. xml_file.write('\n')
    57. if __name__ == "__main__":
    58. yolo = YOLO()
    59. ss = Predict()
    60. #需要修改以下4个量,并且要去VOCdevkit/VOC2007/文件夹下替换训练好的模型best_epoch_weights.pth和voc_classes.txt
    61. pred_class = ["car", "moto", "persons"] # 填入需要检测的类名
    62. file_path = r"D:\AI\4.yolov5-pytorch-main_xml_write\save\image" # 填入测试的图片路径
    63. dir_save_path = r"D:\AI\4.yolov5-pytorch-main_xml_write\save\image_save"# 填入保存的图片路径
    64. xml_path="save\\xml_save\\"# 填入保存的xml文件的路径
    65. ls=os.listdir(file_path)
    66. for item in ls:
    67. img_name=item
    68. xml_name=img_name.split(".")[0]+".xml"
    69. img_names=img_name.split(".")[0]
    70. img_path=os.path.join(file_path,img_name)
    71. save_path=os.path.join(dir_save_path,img_name)
    72. #xml_path=os.path.join(xml_path,xml_name)
    73. pred=ss.a(img_path,save_path,img_name)
    74. img_xml(img_path, xml_path, img_names, pred)

  • 相关阅读:
    四六级听力考试高频词汇分类记忆-日常生活类
    STM32+ MAX30102通过指尖测量心率+血氧饱和度
    Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks
    leetcode刷题集:股票问题六讲【动态规划】(python代码)
    [山东科技大学OJ]2675 Problem C: 回文数字
    android被杀以后fragments缓存重建问题和测试方法
    CentOS7启动SSH服务报错
    基于 Delphi 的前后端分离:之二
    Qt for Android设置应用图标及名称
    el -table 多层级嵌套
  • 原文地址:https://blog.csdn.net/weixin_71719718/article/details/133644961