DroneVehicle数据集转YOLO格式（一标签对应两模态）
由于我需要学习多模态融合，尤其优化融合模块，所以需要将DroneVehicle数据集处理为一标签对应两模态，即在处理过程中，删除双光标注不一致的标注框。
感谢帖子：https://zhuanlan.zhihu.com/p/665126340的鼎立相助！！！
Step1 下载数据集
https://github.com/VisDrone/DroneVehicle
解压得到：
Step2 去除白边
importnumpyasnpimportcv2importosfromtqdmimporttqdmdefcreate_file(output_dir_vi,output_dir_ir):ifnotos.path.exists(output_dir_vi):os.makedirs(output_dir_vi)ifnotos.path.exists(output_dir_ir):os.makedirs(output_dir_ir)print(f'Created folder:({output_dir_vi}); ({output_dir_ir})')defupdate(input_img_path,output_img_path):image=cv2.imread(input_img_path)cropped=image[100:612,100:740]# 裁剪坐标为[y0:y1, x0:x1]cv2.imwrite(output_img_path,cropped)dataset_dir_vi='./data/val/valimg'output_dir_vi='./data/dataset/valimg_cropped'dataset_dir_ir='./data/val/valimgr'output_dir_ir='./data/dataset/valimgr_cropped'# 检查文件夹是否存在，如果不存在则创建create_file(output_dir_vi,output_dir_ir)# 获得需要转化的图片路径并生成目标路径image_filenames_vi=[(os.path.join(dataset_dir_vi,x),os.path.join(output_dir_vi,x))forxinos.listdir(dataset_dir_vi)]image_filenames_ir=[(os.path.join(dataset_dir_ir,x),os.path.join(output_dir_ir,x))forxinos.listdir(dataset_dir_ir)]# 转化所有图片print('Start transforming vision images...')forpathintqdm(image_filenames_vi):update(path[0],path[1])print('Start transforming infrared images...')forpathintqdm(image_filenames_ir):update(path[0],path[1])Step3 XML数据转TXT数据
importosimportxml.etree.ElementTreeasETimportmathimportcv2ascvimportargparsefromtqdmimporttqdm# 图像类别classes=["feright_car","car","truck","bus","van"]# 定义相关地址参数defparse_args():parser=parser=argparse.ArgumentParser(description='polygon')parser.add_argument('--in_xml_vi_dir',default='./test/testlabel',help='可见光 XML 文件地址')parser.add_argument('--in_xml_ir_dir',default='./test/testlabelr',help='红外光 XML 文件地址')parser.add_argument('--out_vi_txt_dir',default='./data_cut/test/testlabel',help='可见光 TXT 文件地址')parser.add_argument('--out_ir_txt_dir',default='./data_cut/test/testlabelr',help='红外光 TXT 文件地址')parser.add_argument('--in_vi_img_dir',default='./data_cut/test/images',help='裁切后的 vi-img 地址')parser.add_argument('--in_ir_img_dir',default='./data_cut/test/images_ir',help='裁切后的 ir-img 地址')args=parser.parse_args()returnargs# 根据 xml 文件中的 name 选择生成的 txt 文件中的 iddefselect_id(name):# 标准化：转小写、去空格（防止大小写或空格问题）name=name.strip().lower()# 定义映射表（支持多种拼写）mapping={# car 类"car":0,# truck 类"truck":1,# bus 类"bus":1,# 注意：有些数据集把 bus 归为 truck？按你需求调整# 如果 bus 是独立类别，保留为 2# van 类"van":3,# freight/feright 相关（统一归为类别 4）"feright":4,"freight":4,"feright car":4,"feright_car":4,"freight car":4,"freight_car":4,"ferightcar":4,"freightcar":4,}ifnameinmapping:returnmapping[name]else:# 打印警告并跳过，避免崩溃print(f"Warning: unknown class name '{name}', skipping this object.")returnNone# YOLO 数据处理defdata_transform(height,width,xmin,ymin,xmax,ymax):# 中心点坐标 x_c,y_cx_c=(xmin+xmax)/2y_c=(ymin+ymax)/2# 中心横坐标与图像宽度比值 x_，中心纵坐标与图像高度比值 y_，bbox 宽度与图像宽度比值 w_，bbox 高度与图像高度比值 h_x_=x_c/width y_=y_c/height w_=(xmax-xmin)/width h_=(ymax-ymin)/heightreturnx_,y_,w_,h_# xml 文件转 txt 文件defxml2txt(in_xml_dir,xml_name,out_txt_dir,in_img_dir):txt_name=xml_name[:-4]+'.txt'# 获取生成的 txt 文件名txt_path=out_txt_dir# 获取生成的 txt 文件保存地址# 判断保存 txt 文件的文件夹是否存在，如果不存在则创建相应文件夹ifnotos.path.exists(txt_path):os.makedirs(txt_path)txt_file=os.path.join(txt_path,txt_name)# 获取 txt 文件地址（保存地址 + 保存名字）img_name=xml_name[:-4]+'.jpg'# 获取图像名字，确保生成的 txt 文件名与图像文件名一致img_path=os.path.join(in_img_dir,img_name)# 获取图像地址img=cv.imread(img_path)# 读取图像信息height,width,_=img.shape# 获取图像高度（height），宽度（width），通道数（_）xml_file=os.path.join(in_xml_dir,xml_name)# 获取 xml 文件地址tree=ET.parse(os.path.join(xml_file))# 使用 ET.parse 方法解析 xml 文件root=tree.getroot()# 使用 getroot 方法获取根目录# 生成对应的 txt 文件withopen(txt_file,"w+",encoding='UTF-8')asout_file:forobjinroot.findall('object'):# 修改部分标注文件中标注不全的 name 文件name=obj.find('name').textifname=='feright_car':name='feright car'else:name=name# 从 xml 文件中提取相关数据信息,并进行删除白边数据操作（白边宽度 100 像素）ifobj.find('polygon'):# 创建空列表用于存放需要处理的数据xmin,xmax,ymin,ymax=[],[],[],[]polygon=obj.find('polygon')# 使用 .find() 方法获取对应 xml 文件中键的键值x1=int(polygon.find('x1').text)-100y1=int(polygon.find('y1').text)-100x2=int(polygon.find('x2').text)-100y2=int(polygon.find('y2').text)-100x3=int(polygon.find('x3').text)-100y3=int(polygon.find('y3').text)-100x4=int(polygon.find('x4').text)-100y4=int(polygon.find('y4').text)-100# 将获取后的数据填入空列表中foriin[x1,x2,x3,x4]:xmin.append(i)xmax.append(i)forjin[y1,y2,y3,y4]:ymin.append(j)ymax.append(j)# 使用 min()、max() 方法获取最大值，最小值xmin=min(xmin)xmax=max(xmax)ymin=min(ymin)ymax=max(ymax)# yolo 格式转换result=data_transform(height,width,xmin,ymin,xmax,ymax)# id 选择result_id=select_id(name)elifobj.find('bndbox'):bndbox=obj.find('bndbox')# 使用 .find() 方法获取对应 xml 文件中键的键值xmin=bndbox.find('xmin').text ymin=bndbox.find('ymin').text xmax=bndbox.find('xmax').text ymax=bndbox.find('ymax').text x1=int(xmin)-100y1=int(ymin)-100x3=int(xmax)-100y3=int(ymax)-100# yolo 格式转换result=data_transform(height,width,x1,y1,x3,y3)# id 选择result_id=select_id(name)# 创建 txt 文件中的数据if0<=result[0]<=1and0<=result[1]<=1and0<=result[2]<=1and0<=result[3]<=1:data=str(result[0])+" "+str(result[1])+" "+str(result[2])+" "+str(result[3])+'\n'data=str(result_id)+" "+data out_file.write(data)else:passif__name__=="__main__":args=parse_args()# 获取命令参数xml_vi_path=args.in_xml_vi_dir# 获取可见光 xml 文件地址xmlFiles_vi=os.listdir(xml_vi_path)# 生成可见光 xml 文件名列表xml_ir_path=args.in_xml_ir_dir# 获取红外 xml 文件地址xmlFiles_ir=os.listdir(xml_ir_path)# 生成红外 xml 文件名列表print('Start transforming vision labels...')foriintqdm(range(0,len(xmlFiles_vi))):xml2txt(args.in_xml_vi_dir,xmlFiles_vi[i],args.out_vi_txt_dir,args.in_vi_img_dir)print('Finish.')print('Start transforming infrared labels...')foriintqdm(range(0,len(xmlFiles_ir))):xml2txt(args.in_xml_ir_dir,xmlFiles_ir[i],args.out_ir_txt_dir,args.in_ir_img_dir)print('Finish.')Step4 滤除不需要的框
该数据集中，对可见光和红外光的标注并不严格对齐。比如存在：仅在一张图片标注（可见光不可见，红外光可见）；位置对齐失败（可见光与红外摄像头规格不同，或标注框出现偏移）。
所以我写了一个脚本，去除多余的框与位置对齐优化，达到可见光和红外光的严格对齐，代码如下：
importosimportcv2importnumpyasnpfrompathlibimportPath# === 配置 ===base_dir="D:\BaiduNetdiskDownload\dronevehicle\data_cut"# 👈 修改为你的真实路径sets=['train','val','test']# 要处理的数据集max_distance=0.02# 中心点距离阈值（相对坐标）# 创建输出目录forsinsets:labels_dir=os.path.join(base_dir,s,"labels")os.makedirs(labels_dir,exist_ok=True)# 读取并解析 YOLO 标签文件defparse_yolo(path):ifnotos.path.exists(path):return[]boxes=[]withopen(path,'r')asf:forlineinf:line=line.strip()ifnotlineorline.startswith('None'):continueparts=line.split()try:cls=int(parts[0])cx,cy,w,h=map(float,parts[1:5])ifcx<0orcy<0orw<=0orh<=0:continueboxes.append((cls,cx,cy,w,h))exceptExceptionase:print(f"⚠️ 解析失败:{path}->{line}| 错误:{e}")returnboxes# 计算两个框的中心点距离defcenter_distance(box1,box2):cx1,cy1=box1[1],box1[2]cx2,cy2=box2[1],box2[2]returnnp.sqrt((cx1-cx2)**2+(cy1-cy2)**2)# 匹配两个列表中的框（基于类别 + 中心点距离）defmatch_boxes(vis_boxes,ir_boxes,max_dist=0.02):matched=[]used_ir=set()forv_boxinvis_boxes:best_match=Nonemin_dist=float('inf')fori,r_boxinenumerate(ir_boxes):ifiinused_ir:continueifv_box[0]!=r_box[0]:# 类别不同，跳过continuedist=center_distance(v_box,r_box)ifdist<max_distanddist<min_dist:min_dist=dist best_match=(i,r_box)ifbest_match:idx,r_box=best_match used_ir.add(idx)# 取平均坐标avg_cls=v_box[0]avg_cx=(v_box[1]+r_box[1])/2avg_cy=(v_box[2]+r_box[2])/2avg_w=(v_box[3]+r_box[3])/2avg_h=(v_box[4]+r_box[4])/2matched.append((avg_cls,avg_cx,avg_cy,avg_w,avg_h))returnmatched# 处理每个数据集fordatasetinsets:# 构造路径：trainlabel, trainlabelr 等label_dir=os.path.join(base_dir,dataset,f"{dataset}label")# 如 trainlabellabel_r_dir=os.path.join(base_dir,dataset,f"{dataset}labelr")# 如 trainlabelroutput_label_dir=os.path.join(base_dir,dataset,"labels")# 获取所有样本 ID（以 .txt 结尾的文件名，去掉后缀）sample_ids=set()forpathinPath(label_dir).glob("*.txt"):sample_ids.add(os.path.basename(path).replace(".txt",""))forpathinPath(label_r_dir).glob("*.txt"):sample_ids.add(os.path.basename(path).replace(".txt",""))print(f"🔍 正在处理{dataset}，共{len(sample_ids)}个样本...")forsample_idinsample_ids:vis_path=os.path.join(label_dir,f"{sample_id}.txt")ir_path=os.path.join(label_r_dir,f"{sample_id}.txt")# 读取两个标签vis_boxes=parse_yolo(vis_path)ir_boxes=parse_yolo(ir_path)# 如果任一方为空，跳过iflen(vis_boxes)==0orlen(ir_boxes)==0:print(f"⚠️{sample_id}: 一方无标注，跳过")continue# 匹配框matched_boxes=match_boxes(vis_boxes,ir_boxes,max_distance)# 如果没有匹配，跳过iflen(matched_boxes)==0:print(f"⚠️{sample_id}: 无匹配框，跳过")continue# 写入输出文件out_path=os.path.join(output_label_dir,f"{sample_id}.txt")withopen(out_path,'w')asf:forboxinmatched_boxes:f.write(f"{box[0]}{box[1]:.6f}{box[2]:.6f}{box[3]:.6f}{box[4]:.6f}\n")print(f"✅{sample_id}: 生成{len(matched_boxes)}个对齐框 →{out_path}")print("\n🎉 所有数据集处理完成！")Step5 可视化
importosimportcv2importnumpyasnpimportrandomfrompathlibimportPath# === 配置 ===base_dir="D:\BaiduNetdiskDownload\dronevehicle\data_cut"# 👈 修改为你的真实路径sets=['train','val','test']# 图像后缀img_exts=['.jpg','.jpeg','.png']ir_exts=['.png','.jpg']# 红外可能为 pngdefload_image(path,exts):forextinexts:p=path+extifos.path.exists(p):returncv2.imread(p)print(f"❌ 图像未找到:{path}")returnNonedefparse_labels(label_path):boxes=[]ifnotos.path.exists(label_path):print(f"⚠️ 标签文件不存在:{label_path}")returnboxeswithopen(label_path,'r')asf:forlineinf:line=line.strip()ifnotlineorline.startswith('None'):continueparts=line.split()try:cls=int(parts[0])cx,cy,w,h=map(float,parts[1:5])ifcx<0orcy<0orw<=0orh<=0:continueboxes.append((cls,cx,cy,w,h))exceptExceptionase:print(f"⚠️ 解析失败:{line}| 错误:{e}")returnboxesdefdraw_boxes(img,boxes):"""在图像上绘制 YOLO 格式的框"""h,w=img.shape[:2]forcls,cx,cy,w_box,h_boxinboxes:x1=int((cx-w_box/2)*w)y1=int((cy-h_box/2)*h)x2=int((cx+w_box/2)*w)y2=int((cy+h_box/2)*h)color=(0,255,0)ifcls==0else(0,0,255)cv2.rectangle(img,(x1,y1),(x2,y2),color,2)cv2.putText(img,str(cls),(x1,y1-10),cv2.FONT_HERSHEY_SIMPLEX,0.6,color,2)returnimgdefget_all_samples(base_dir,dataset):"""获取指定数据集的所有样本 ID"""label_dir=os.path.join(base_dir,dataset,"labels")ifnotos.path.exists(label_dir):return[]return[p.stemforpinPath(label_dir).glob("*.txt")]# 主程序defmain():all_samples=[]forsinsets:samples=get_all_samples(base_dir,s)all_samples.extend([(s,sid)forsidinsamples])ifnotall_samples:print("❌ 没有找到任何样本，请检查 labels/ 目录！")returnprint(f"🎉 共找到{len(all_samples)}个样本，开始可视化...")random.shuffle(all_samples)current_idx=0total=len(all_samples)whileTrue:dataset,sample_id=all_samples[current_idx]print(f"\n➡️ 当前样本:{dataset}/{sample_id}")# 构造路径image_path=os.path.join(base_dir,dataset,"images",sample_id)ir_path=os.path.join(base_dir,dataset,"images_ir",sample_id)label_path=os.path.join(base_dir,dataset,"labels",f"{sample_id}.txt")# 加载图像vis_img=load_image(image_path,img_exts)ir_img=load_image(ir_path,ir_exts)ifvis_imgisNoneorir_imgisNone:print(f"❌ 图像加载失败，跳过{sample_id}")current_idx+=1ifcurrent_idx>=total:current_idx=0continue# 加载标签boxes=parse_labels(label_path)# 绘制框vis_with_boxes=draw_boxes(vis_img.copy(),boxes)ir_with_boxes=draw_boxes(ir_img.copy(),boxes)# 显示两幅图cv2.imshow("Visible Light",vis_with_boxes)cv2.imshow("Infrared",ir_with_boxes)# 等待按键key=cv2.waitKey(0)&0xFFifkey==ord('q')orkey==27:# q 或 ESC 退出breakelifkey==ord('n'):# n 下一个current_idx+=1else:current_idx+=1# 默认按任意键下一页# 循环ifcurrent_idx>=total:current_idx=0cv2.destroyAllWindows()if__name__=="__main__":main()用于观察是否达到严格对齐，效果不好可以修改：
max_distance = 0.02 # 中心点距离阈值（相对坐标）
