河北建设厅网站首页,美图秀秀网页版,网站建设提供资料表,建站及推广yolov5优化模型时#xff0c;一般需要继续标注一些检测错误的图片#xff0c;将其标为xml数据。以下是根据训练好的模型自动标注xml数据的python代码#xff1a;
注意#xff1a;代码中包含了本人的yolov5的测试过程#xff0c;测试过程可以自己根据yolov5的测试文件自行…yolov5优化模型时一般需要继续标注一些检测错误的图片将其标为xml数据。以下是根据训练好的模型自动标注xml数据的python代码
注意代码中包含了本人的yolov5的测试过程测试过程可以自己根据yolov5的测试文件自行修改只是测试返回的类格式为
[[water,[15,20,30,40]],[red,[12,13,14,15]]]
二维数组表示测试的类为water和red,其中后面的数字表示类的坐标[top,left,bottom,right],表示上、左、下、右4个坐标。 import os
import cv2
from PIL import Imagefrom yolo import YOLO#1.预测类获得字符串
class Predict():def a(self, img_path,save_path,img_name):image Image.open(img_path)r_image, pred yolo.detect_image(image, pred_class, img_name)if not os.path.exists(dir_save_path):os.makedirs(dir_save_path)r_image.save(save_path, quality95, subsampling0)return pred#2.写入xml文件
def img_xml(img_path,xml_path,img_name,pred):if len(pred) ! 0:#1.读取图片(xml需要写入图片的长宽高)img cv2.imread(img_path)#2.写入xml文件#(1)写入文件头部files_pathimg_path.split(\\)[-2]print(..:,files_path)xml_file open((xml_path img_name .xml), w)xml_file.write(annotation\n)xml_file.write( folder files_path /folder\n)xml_file.write( filename img_name .jpg /filename\n)xml_file.write( path img_path /path\n)xml_file.write( source\n)xml_file.write( databaseUnknown/database\n)xml_file.write( /source\n)#2写入图片的长宽高信息xml_file.write( size\n)xml_file.write( widthstr(img.shape[1])/width\n)xml_file.write( height str(img.shape[0]) /height\n)xml_file.write( depth str(img.shape[2]) /depth\n)xml_file.write( /size\n)xml_file.write( segmented0/segmented\n)#3.写入字符串信息:[[water,[15,20,30,40]],[red,[12,13,14,15]]]#if len(shuzu)!0:for item in pred:xml_file.write( object\n)xml_file.write( name str(item[0]) /name\n)xml_file.write( poseUnspecified/pose\n)xml_file.write( truncated0/truncated\n)xml_file.write( difficult0/difficult\n)xml_file.write( bndbox\n)#写入字符串信息#[top, left, bottom, right]xml_file.write( xmin str(item[1][1]) /xmin\n)xml_file.write( ymin str(item[1][0]) /ymin\n)xml_file.write( xmax str(item[1][3]) /xmax\n)xml_file.write( ymax str(item[1][2]) /ymax\n)xml_file.write( /bndbox\n)xml_file.write( /object\n)xml_file.write(/annotation\n)if __name__ __main__:yolo YOLO()ss Predict()#需要修改以下4个量并且要去VOCdevkit/VOC2007/文件夹下替换训练好的模型best_epoch_weights.pth和voc_classes.txtpred_class [car, moto, persons] # 填入需要检测的类名file_path rD:\AI\4.yolov5-pytorch-main_xml_write\save\image # 填入测试的图片路径dir_save_path rD:\AI\4.yolov5-pytorch-main_xml_write\save\image_save# 填入保存的图片路径xml_pathsave\\xml_save\\# 填入保存的xml文件的路径lsos.listdir(file_path)for item in ls:img_nameitemxml_nameimg_name.split(.)[0].xmlimg_namesimg_name.split(.)[0]img_pathos.path.join(file_path,img_name)save_pathos.path.join(dir_save_path,img_name)#xml_pathos.path.join(xml_path,xml_name)predss.a(img_path,save_path,img_name)img_xml(img_path, xml_path, img_names, pred)
