針對的是多邊形(<polygon>
)來描述對象的邊界,而不是傳統的矩形框(<bndbox>
)
import xml.etree.ElementTree as ET
import os
from pathlib import Path# 解析VOC格式的XML文件,提取目標框的標簽和坐標
def parse_voc_xml(xml_file):tree = ET.parse(xml_file)root = tree.getroot()# 獲取圖像的寬度和高度size = root.find('size')img_width = int(size.find('width').text)img_height = int(size.find('height').text)bboxes = []for obj in root.findall('object'):label = obj.find('name').text # 獲取目標的標簽polygon = obj.find('polygon') # 獲取多邊形的坐標# 獲取多邊形的所有頂點坐標x_coords = []y_coords = []for i in range(1, 5): # 假設最多4個點(也可以擴展)x_coords.append(int(polygon.find(f'x{i}').text))y_coords.append(int(polygon.find(f'y{i}').text))# 計算邊界框(最小矩形框)的坐標xmin = min(x_coords)xmax = max(x_coords)ymin = min(y_coords)ymax = max(y_coords)# 將 VOC 格式的坐標轉換為 YOLO 格式x_center = (xmin + xmax) / 2y_center = (ymin + ymax) / 2width = xmax - xminheight = ymax - ymin# 歸一化坐標x_center /= img_widthy_center /= img_heightwidth /= img_widthheight /= img_heightbboxes.append((label, x_center, y_center, width, height))return bboxes, img_width, img_height# 將標簽保存為 YOLO 格式
def save_yolo_labels(image_filename, bboxes, output_dir):# YOLO 標簽文件的路徑txt_filename = os.path.join(output_dir, Path(image_filename).stem + '.txt')with open(txt_filename, 'w') as f:for label, x_center, y_center, width, height in bboxes:# YOLO 格式為:<class_id> <x_center> <y_center> <width> <height># 假設標簽的類別是數字(如果是字符串,則需要將類別轉換為數字映射)f.write(f"{label} {x_center} {y_center} {width} {height}\n")# 主函數:將 VOC XML 文件轉換為 YOLO 格式
def convert_voc_to_yolo(voc_xml_dir, output_dir, class_mapping):# 如果輸出目錄不存在,則創建if not os.path.exists(output_dir):os.makedirs(output_dir)# 遍歷所有 XML 文件for xml_file in Path(voc_xml_dir).glob('*.xml'):print(f"Processing {xml_file}")# 解析 VOC XML 文件bboxes, img_width, img_height = parse_voc_xml(xml_file)# 將標簽轉化為 YOLO 格式,且需要使用數字類別而不是字符串bboxes = [(class_mapping[label], x_center, y_center, width, height) for label, x_center, y_center, width, height in bboxes]# 保存 YOLO 格式的標簽save_yolo_labels(xml_file.stem + '.jpg', bboxes, output_dir) # 假設圖像文件與XML文件同名,后綴為 .jpg# 主程序入口
if __name__ == "__main__":# 輸入目錄:包含VOC XML標注文件的目錄voc_xml_dir = './original/labels' # 替換為 VOC XML 文件目錄路徑# 輸出目錄:YOLO 格式標簽文件的存儲目錄output_dir = './original/save' # 替換為 YOLO 標簽文件存儲目錄路徑# 類別映射:將VOC中的標簽名映射為數字類別class_mapping = {'car': 0, 'truck': 1,'bus': 2,'motorcycle': 3,'van': 4,'freight_car': 5,}# 轉換 VOC 格式標簽為 YOLO 格式convert_voc_to_yolo(voc_xml_dir, output_dir, class_mapping)
# 'car', 'truck', 'bus', 'motorcycle', 'van', 'freight_car'