Anti-UAV 反无人机大赛数据集处理脚本(3)YOLO结果转提交格式

import os
import json
from PIL import Image

def convert_yolo_to_xywh(yolo_box, img_w, img_h):
    """将YOLO格式转换为左上角坐标+宽高格式（保留浮点精度）"""
    cx, cy, w, h = map(float, yolo_box)
    
    # 计算绝对坐标（保持浮点运算）
    cx_abs = cx * img_w
    cy_abs = cy * img_h
    w_abs = w * img_w
    h_abs = h * img_h
    
    # 计算左上角坐标
    x = cx_abs - w_abs / 2
    y = cy_abs - h_abs / 2
    
    # 边界检查（使用浮点比较）
    x = max(0.0, min(x, img_w - 1.0))
    y = max(0.0, min(y, img_h - 1.0))
    w = max(0.0, min(w_abs, img_w - x))
    h = max(0.0, min(h_abs, img_h - y))
    
    return [x, y, w, h]

def process_sequence(image_dir, label_dir):
    """处理单个序列（精确浮点版本）"""
    results = []
    
    images = sorted(
        [f for f in os.listdir(image_dir) if f.lower().endswith(('.jpg', '.png', '.jpeg'))],
        key=lambda x: int(os.path.splitext(x)[0])
    )
    
    for img_file in images:
        img_path = os.path.join(image_dir, img_file)
        base_name = os.path.splitext(img_file)[0]
        label_path = os.path.join(label_dir, f"{base_name}.txt")
        
        try:
            with Image.open(img_path) as img:
                img_w, img_h = img.size
            
            if not os.path.exists(label_path):
                results.append([0.0, 0.0, 0.0, 0.0])  # 保持浮点格式
                continue
                
            with open(label_path, 'r') as f:
                lines = f.readlines()
            
            if not lines:
                results.append([0.0, 0.0, 0.0, 0.0])
                continue
                
            parts = lines[0].strip().split()
            if len(parts) < 5:
                results.append([0.0, 0.0, 0.0, 0.0])
                continue
                
            try:
                converted = convert_yolo_to_xywh(parts[1:5], img_w, img_h)
                # 保留原始浮点精度
                results.append([
                    float(f"{converted[0]:.15f}"),
                    float(f"{converted[1]:.15f}"),
                    float(f"{converted[2]:.15f}"),
                    float(f"{converted[3]:.15f}")
                ])
            except:
                results.append([0, 0, 0, 0])
                
        except Exception as e:
            print(f"Error processing {img_file}: {str(e)}")
            results.append([0, 0, 0, 0])
    
    return results

def main(root_dir, image_root, output_dir):
    """主处理函数"""
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 遍历所有子文件夹
    for seq_name in os.listdir(root_dir):
        seq_path = os.path.join(root_dir, seq_name)
        if not os.path.isdir(seq_path):
            continue
            
        # 对应图片目录
        img_seq_dir = os.path.join(image_root, seq_name)
        if not os.path.exists(img_seq_dir):
            print(f"Warning: Missing image folder for {seq_name}")
            continue
            
        # 标签目录
        label_dir = os.path.join(seq_path, 'labels')
        if not os.path.exists(label_dir):
            print(f"Warning: Missing labels folder for {seq_name}")
            continue
            
        # 处理序列
        print(f"Processing {seq_name}...")
        res = process_sequence(img_seq_dir, label_dir)
        
        # 写入结果文件
        output_path = os.path.join(output_dir, f"{seq_name}.txt")
        with open(output_path, 'w') as f:
            json.dump({"res": res}, f)
            
if __name__ == "__main__":
    # 配置路径
    ROOT_DIR = "/kaijiang/xy/ultralytics-yolo11-20250301/ultralytics-yolo11-main/runs/detect"        # 包含子文件夹的根目录
    IMAGE_ROOT = "/kaijiang/xy/ultralytics-yolo11-20250301/ultralytics-yolo11-main/dataset/antiuavtest/track2_test"    # 包含同名子文件夹的图片目录
    OUTPUT_DIR = "/kaijiang/xy/ultralytics-yolo11-20250301/ultralytics-yolo11-main/myresult"   # 输出目录
    
    main(ROOT_DIR, IMAGE_ROOT, OUTPUT_DIR)