三维重构终版

OpenCV/convert_to_voc.py

@@ -0,0 +1,192 @@
+import os
+import json
+import xml.etree.ElementTree as ET
+from xml.dom import minidom
+from tqdm import tqdm
+import re  # 引入正则表达式库
+
+# --- 配置参数 ---
+# 1. 原始JSON文件所在的文件夹路径
+json_folder = '../label/up_json'  # 示例路径，请修改为您的JSON文件夹
+
+# 2. 原始图片文件所在的文件夹路径 (用于获取图片尺寸)
+image_folder = '../label/up'  # 示例路径，请修改为您的图片文件夹
+
+# 3. 生成的XML文件要保存的文件夹路径
+output_xml_folder = '../label/up_xml'
+
+# 4. 您要检测的目标类别名称 (对应 label "3")
+class_name_for_label_3 = "Space weld workpiece"  # 这是您XML示例中的名称
+
+# 5. 分组的大小
+group_size = 5
+
+
+# --- 配置结束 ---
+
+
+def create_xml_annotation(image_info, objects_info):
+    """
+    根据传入的信息生成XML树对象
+    :param image_info: 包含图片文件名、尺寸等信息的字典
+    :param objects_info: 包含多个物体信息的列表，每个物体是一个字典
+    :return: XML ElementTree对象
+    """
+    # 创建根节点
+    annotation = ET.Element('annotation')
+
+    # 子节点 - folder
+    folder = ET.SubElement(annotation, 'folder')
+    folder.text = 'JPEGImages'
+
+    # 子节点 - filename
+    filename_node = ET.SubElement(annotation, 'filename')
+    filename_node.text = image_info['filename']
+
+    # 子节点 - path (路径通常不那么重要，但最好有一个)
+    path = ET.SubElement(annotation, 'path')
+    # 路径指向JPEGImages文件夹
+    image_path_in_voc = os.path.join('..', 'JPEGImages', image_info['filename'])
+    path.text = image_path_in_voc
+
+    # 子节点 - source
+    source = ET.SubElement(annotation, 'source')
+    database = ET.SubElement(source, 'database')
+    database.text = 'Unknown'
+
+    # 子节点 - size
+    size = ET.SubElement(annotation, 'size')
+    width = ET.SubElement(size, 'width')
+    width.text = str(image_info['width'])
+    height = ET.SubElement(size, 'height')
+    height.text = str(image_info['height'])
+    depth = ET.SubElement(size, 'depth')
+    depth.text = str(image_info.get('depth', 3))
+
+    # 子节点 - segmented
+    segmented = ET.SubElement(annotation, 'segmented')
+    segmented.text = '0'
+
+    # 为每个物体添加 object 节点
+    for obj in objects_info:
+        object_node = ET.SubElement(annotation, 'object')
+        name = ET.SubElement(object_node, 'name')
+        name.text = obj['name']
+        pose = ET.SubElement(object_node, 'pose')
+        pose.text = 'Unspecified'
+        truncated = ET.SubElement(object_node, 'truncated')
+        truncated.text = '0'
+        difficult = ET.SubElement(object_node, 'difficult')
+        difficult.text = '0'
+        bndbox = ET.SubElement(object_node, 'bndbox')
+        xmin = ET.SubElement(bndbox, 'xmin')
+        xmin.text = str(int(obj['xmin']))
+        ymin = ET.SubElement(bndbox, 'ymin')
+        ymin.text = str(int(obj['ymin']))
+        xmax = ET.SubElement(bndbox, 'xmax')
+        xmax.text = str(int(obj['xmax']))
+        ymax = ET.SubElement(bndbox, 'ymax')
+        ymax.text = str(int(obj['ymax']))
+
+    return annotation
+
+
+def prettify_xml(elem):
+    """
+    格式化XML输出，使其更易读
+    """
+    rough_string = ET.tostring(elem, 'utf-8')
+    reparsed = minidom.parseString(rough_string)
+    return reparsed.toprettyxml(indent="  ")
+
+
+def main():
+    if not os.path.exists(output_xml_folder):
+        os.makedirs(output_xml_folder)
+        print(f"创建输出文件夹: {output_xml_folder}")
+
+    json_files = sorted([f for f in os.listdir(json_folder) if f.endswith('.json')])
+
+    print(f"找到 {len(json_files)} 个JSON文件，开始转换...")
+
+    for json_file in tqdm(json_files, desc="处理JSON文件"):
+        base_name = os.path.splitext(json_file)[0]
+
+        # 使用正则表达式匹配前缀和数字
+        match = re.match(r'([a-zA-Z]+)(\d+)', base_name)
+
+        # 1. 检查当前文件是否是一个分组的起始文件
+        is_group_start_file = False
+        if match:
+            num = int(match.group(2))
+            # 如果数字是 1, 6, 11, ... 这样的，就认为是起始文件
+            if (num - 1) % group_size == 0:
+                is_group_start_file = True
+        else:
+            # 如果文件名不符合 l1, r5 这种格式，我们认为它是“普通”文件，自己就是一个组
+            is_group_start_file = True
+
+        if not is_group_start_file:
+            # 如果不是起始文件（如l2, l3...），则跳过，因为它的标注已由l1处理
+            continue
+
+        # --- 是起始文件，处理这个分组 ---
+        json_path = os.path.join(json_folder, json_file)
+
+        with open(json_path, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+
+        # 2. 从起始文件中提取所有符合条件的标注对象
+        objects_to_write = []
+        for shape in data.get('shapes', []):
+            if shape.get('label') == '1' and shape.get('shape_type') == 'rectangle':
+                points = shape.get('points', [])
+                if len(points) == 2:
+                    x_coords = sorted([p[0] for p in points])
+                    y_coords = sorted([p[1] for p in points])
+                    objects_to_write.append({
+                        'name': class_name_for_label_3,
+                        'xmin': x_coords[0], 'ymin': y_coords[0],
+                        'xmax': x_coords[1], 'ymax': y_coords[1],
+                    })
+
+        if not objects_to_write:
+            continue
+
+        # 3. 确定该标注要应用到哪些图片上
+        target_image_names = []
+        if match:
+            # 文件名符合 l1, r6 等格式
+            prefix = match.group(1)
+            start_num = int(match.group(2))
+            for i in range(group_size):
+                # 假设图片格式为 .jpg
+                target_image_names.append(f"{prefix}{start_num + i}.jpg")
+        else:
+            # 普通文件，只应用到同名文件
+            # 假设图片格式为 .jpg
+            target_image_names.append(f"{base_name}.jpg")
+
+        # 4. 为分组内的每个目标图片生成XML文件
+        for image_name in target_image_names:
+            image_path = os.path.join(image_folder, image_name)
+            if not os.path.exists(image_path):
+                print(f"\n警告：找不到图片 '{image_name}'，跳过生成其XML文件。")
+                continue
+
+            # 使用JSON中的尺寸信息
+            image_info = {'filename': image_name, 'width': data['imageWidth'], 'height': data['imageHeight']}
+
+            xml_tree = create_xml_annotation(image_info, objects_to_write)
+            xml_string = prettify_xml(xml_tree)
+            xml_filename = os.path.splitext(image_name)[0] + '.xml'
+            output_path = os.path.join(output_xml_folder, xml_filename)
+
+            with open(output_path, 'w', encoding='utf-8') as f:
+                f.write(xml_string)
+
+    print("转换完成！所有XML文件已保存在: ", output_xml_folder)
+
+
+if __name__ == '__main__':
+    main()