资讯动态

基于YOLO识别水表框

发布时间:2026/9/5 18:05:15 来源:尧图企业网站定制
一、将.xml文件当中相关的wmwm1标签的具体信息转换为.txt文件import os import xml.etree.ElementTree as ET import math # 类别列表必须与XML中的name标签完全一致 cls_list [ wm, wm1, wm2, wm3, wm4, wm5, wm6, wm7, wm8, wm9, wm10, wm11, wm12, wm13, wm14, wm15, wm16, wm17, wm18, wm19, wm20, wm21 ] def convert_rotated_rect(cx, cy, w, h, angle, img_w, img_h): 将旋转矩形转换为四个角点坐标归一化 points [] for dx, dy in [(-w / 2, -h / 2), (w / 2, -h / 2), (w / 2, h / 2), (-w / 2, h / 2)]: x cx dx * math.cos(angle) - dy * math.sin(angle) y cy dx * math.sin(angle) dy * math.cos(angle) # 坐标限幅并归一化 x_norm max(0, min(x, img_w)) / img_w y_norm max(0, min(y, img_h)) / img_h points.extend([x_norm, y_norm]) return points def process_xml(xml_path, output_dir): 处理单个XML文件 tree ET.parse(xml_path) root tree.getroot() # 获取图像尺寸 size root.find(size) img_w float(size.find(width).text) img_h float(size.find(height).text) # 准备输出文件路径 filename os.path.splitext(os.path.basename(xml_path))[0] txt_path os.path.join(output_dir, f{filename}.txt) with open(txt_path, w) as f: for obj in root.findall(object): try: # 解析类别 cls_name obj.find(name).text.strip() if cls_name not in cls_list: continue # 解析旋转框参数 robndbox obj.find(robndbox) cx float(robndbox.find(cx).text) cy float(robndbox.find(cy).text) w float(robndbox.find(w).text) h float(robndbox.find(h).text) angle float(robndbox.find(angle).text) # 转换坐标并归一化 points convert_rotated_rect(cx, cy, w, h, angle, img_w, img_h) # 获取类别索引 cls_index cls_list.index(cls_name) # 写入DOTA格式仅保留cls_index和8个坐标 line f{cls_index} .join(f{p:.6f} for p in points) \n f.write(line) except Exception as e: print(f处理对象时出错{str(e)}) def batch_convert(xml_dir, output_dir): 批量处理XML目录 os.makedirs(output_dir, exist_okTrue) for filename in os.listdir(xml_dir): if filename.endswith(.xml): xml_path os.path.join(xml_dir, filename) process_xml(xml_path, output_dir) print(f已处理{filename}) if __name__ __main__: # 配置路径使用绝对路径 XML_DIR os.path.abspath(r.\labels_xml) # XML文件所在目录 OUTPUT_DIR os.path.abspath(r.\labels_txt) # 输出目录 # 执行转换 batch_convert(XML_DIR, OUTPUT_DIR) print(f转换完成输出目录{OUTPUT_DIR})其中各个含义代表了x1 y1 x2 y2 x3 y3 x4 y4 class_name difficulty运行以上代码转换.txt文件的格式二、划分数据集import os import random import shutil random.seed(42) 该脚本用于将给定的数据集分割成训练集和测试集。 数据集应包含图像和对应的标注文件。 脚本会按照90%训练集和10%测试集的比例进行分割并将图像和标注文件分别复制到相应的文件夹中。 # 设置数据集文件夹路径和输出文件夹路径 data_folder ./dataset ori_img_folder ./images ori_label_folder ./label img_folder ./dataset/images label_folder ./dataset/labels img_train img_folder /train img_val img_folder /val label_train label_folder /train_original label_val label_folder /val_original if not os.path.exists(img_train): # 检查目录是否存在 os.mkdir(img_train) if not os.path.exists(img_val): # 检查目录是否存在 os.mkdir(img_val) if not os.path.exists(label_train): # 检查目录是否存在 os.mkdir(label_train) if not os.path.exists(label_val): # 检查目录是否存在 os.mkdir(label_val) # 计算每个子集的大小 # 总文件数乘以0.9得到训练集大小其余为测试集大小 total_files len(os.listdir(ori_img_folder)) train_size int(total_files * 0.9) test_size int(total_files - train_size) # 获取所有图像文件的文件名列表并进行随机打乱 image_files os.listdir(ori_img_folder) random.shuffle(image_files) # 复制图像和标注文件到相应的子集文件夹中 # 枚举每个图像文件根据索引决定复制到训练集还是测试集文件夹 for i, image_file in enumerate(image_files): base_file_name os.path.splitext(image_file)[0] # 获取文件名不包括扩展名 image_path os.path.join(ori_img_folder, image_file) label_path os.path.join(ori_label_folder, base_file_name .txt) # 根据索引判断文件应复制到训练集还是测试集 if i train_size: shutil.copy(image_path, os.path.join(img_folder, train)) # 复制图像到训练集 shutil.copy(label_path, os.path.join(label_folder, train_original)) # 复制标注到训练集 else: shutil.copy(image_path, os.path.join(img_folder, val)) # 复制图像到测试集 shutil.copy(label_path, os.path.join(label_folder, val_original)) # 复制标注到测试集from ultralytics.data.converter import convert_dota_to_yolo_obb convert_dota_to_yolo_obb(./dataset)四、训练模型path: /home/dell/Disk4TB/xxc/model1/dataset train: images/train val: images/val # 类别配置 nc: 2 names: [wm,wm1] # OBB 专用参数 obb_format: xyxyxyxy # 8参数旋转框格式import torch from ultralytics import YOLO def main(): modelYOLO(yolov8n-obb.pt) device 0 if torch.cuda.is_available() else cpu # 使用第一个 GPU若无 GPU 则使用 CPU model.train(data./dataset.yaml, epochs200, imgsz640, batch16, workers0, devicedevice) if __name__ __main__: main()

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价