generate_final_tables.py

import os
import torch
from ultralytics import YOLO

# Fix for OpenMP error
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

def benchmark_model_custom(data_dir, model_name, weights_path='yolov10m.pt'):
    print(f"\n🚀 Benchmarking: {model_name}")
    print(f"📂 Data: {data_dir}")
    
    # 1. Create a temporary YAML for this dataset if needed
    # We can reuse the existing train/val split logic by just mapping the paths
    # But since we generated 'data/rtts_original_val.txt' which contains relative paths like 'data/RTTS_Original_Standard/images/...'
    # We need to create a similar txt for the new dataset 'data/RTTS_Dehazed_Original_AOD'
    
    # Let's generate the val.txt first
    val_txt_path = f"data/rtts_{model_name.lower().replace(' ', '_')}_val.txt"
    
    # Read original val txt to get file names
    with open("data/rtts_original_val.txt", 'r') as f:
        orig_paths = f.readlines()
        
    # Replace the root directory
    # Original: data/RTTS_Original_Standard/images/name.png
    # New:      data/RTTS_Dehazed_Original_AOD/images/name.png
    
    # Check what the actual path in orig_paths looks like
    # It might be relative or absolute.
    
    new_paths = []
    for p in orig_paths:
        p = p.strip()
        if not p: continue
        
        # We know the structure is .../images/filename.png
        filename = os.path.basename(p)
        
        # Construct new path
        new_p = os.path.join(data_dir, "images", filename)
        new_paths.append(new_p)
        
    with open(val_txt_path, 'w') as f:
        f.write('\n'.join(new_paths))
        
    print(f"Created val list: {val_txt_path}")
    
    # Create YAML
    yaml_path = f"data/rtts_{model_name.lower().replace(' ', '_')}.yaml"
    with open("data/rtts_original_train.yaml", 'r') as f:
        base_yaml = f.read()
        
    # Replace val path
    # We can just write a new simple yaml
    yaml_content = f"""
path: .
train: {val_txt_path} # We don't need train for benchmark, but YOLO needs the key
val: {val_txt_path}

names:
  0: person
  1: bicycle
  2: car
  3: motorcycle
  4: airplane
  5: bus
  6: train
  7: truck
  8: boat
  9: traffic light
  10: fire hydrant
  11: stop sign
  12: parking meter
  13: bench
  14: bird
  15: cat
  16: dog
  17: horse
  18: sheep
  19: cow
  20: elephant
  21: bear
  22: zebra
  23: giraffe
  24: backpack
  25: umbrella
  26: handbag
  27: tie
  28: suitcase
  29: frisbee
  30: skis
  31: snowboard
  32: sports ball
  33: kite
  34: baseball bat
  35: baseball glove
  36: skateboard
  37: surfboard
  38: tennis racket
  39: bottle
  40: wine glass
  41: cup
  42: fork
  43: knife
  44: spoon
  45: bowl
  46: banana
  47: apple
  48: sandwich
  49: orange
  50: broccoli
  51: carrot
  52: hot dog
  53: pizza
  54: donut
  55: cake
  56: chair
  57: couch
  58: potted plant
  59: bed
  60: dining table
  61: toilet
  62: tv
  63: laptop
  64: mouse
  65: remote
  66: keyboard
  67: cell phone
  68: microwave
  69: oven
  70: toaster
  71: sink
  72: refrigerator
  73: book
  74: clock
  75: vase
  76: scissors
  77: teddy bear
  78: hair drier
  79: toothbrush
"""
    with open(yaml_path, 'w') as f:
        f.write(yaml_content)
        
    # Run Validation
    try:
        model = YOLO(weights_path)
        metrics = model.val(
            data=yaml_path,
            batch=16,
            imgsz=640,
            conf=0.001,
            iou=0.6,
            split='val',
            project='runs/benchmark_tables',
            name=model_name,
            classes=[2, 3, 5], # Car, Motor, Bus
            verbose=False,
            plots=False
        )
        
        ap50 = metrics.box.ap50
        indices = metrics.ap_class_index
        ap_dict = {idx: ap for idx, ap in zip(indices, ap50)}
        
        return {
            "Car": ap_dict.get(2, 0.0),
            "Motorcycle": ap_dict.get(3, 0.0),
            "Bus": ap_dict.get(5, 0.0),
            "mAP": (ap_dict.get(2,0)+ap_dict.get(3,0)+ap_dict.get(5,0))/3
        }
        
    except Exception as e:
        print(f"❌ Error: {e}")
        return None

def generate_final_tables():
    results = {}
    
    # 1. YOLOv10 (Baseline on Foggy) - Pretrained
    # We already have this data, but running it ensures consistency
    # Or we can just use the known values to save time if confirmed.
    # User said "re-run if needed". Let's be safe and quick-run or reuse.
    # Given we just ran it in previous step (71.94%), we can hardcode OR re-run.
    # Let's re-run "YOLO + Original AOD" since that's NEW.
    # Others we can reuse from benchmark_extended.py logs if we trust them.
    
    # Let's run the MISSING piece: YOLO + Original AOD (Pretrained)
    print("--- 1. Benchmarking YOLO + Original AOD ---")
    res_aod = benchmark_model_custom(
        "data/RTTS_Dehazed_Original_AOD", 
        "yolo_aod_original",
        "yolov10m.pt"
    )
    if res_aod:
        results["YOLO+AOD"] = res_aod
        
    # We also need "YOLO+AOD-Advanced" (Pretrained)
    # We have this from previous run: Car 71.80%. 
    # But let's verify if user wants "Pretrained" or "Finetuned" for Table 1.
    # Table 1 implies "Comparison of Dehazing Algorithms" usually without fine-tuning YOLO for each.
    # So we use Pretrained YOLO for all rows in Table 1.
    
    print("--- 2. Benchmarking YOLO + AOD-Advanced (Pretrained) ---")
    res_adv = benchmark_model_custom(
        "data/RTTS_Dehazed_Advanced_Perceptual",
        "yolo_aod_advanced",
        "yolov10m.pt"
    )
    if res_adv:
        results["YOLO+Advanced"] = res_adv
        
    # Baseline YOLO (Foggy)
    print("--- 3. Benchmarking YOLO (Foggy) ---")
    res_base = benchmark_model_custom(
        "data/RTTS_Original_Standard",
        "yolo_foggy",
        "yolov10m.pt"
    )
    if res_base:
        results["YOLO"] = res_base

    # --- 4. Benchmarking Improved YOLO (Ours Finetuned) ---
    print("\n🚀 Benchmarking Improved YOLO (Ours Finetuned)")
    
    # Path to best weights from extended training
    weights_ours = "runs/train_extended/yolov10m_perceptual_extended/weights/best.pt"
    if not os.path.exists(weights_ours):
         # Fallback to interrupted run
         weights_ours = "runs/detect/runs/train_comparison/yolov10m_rtts_perceptual_train2/weights/best.pt"
    
    if os.path.exists(weights_ours):
        print(f"Using weights: {weights_ours}")
        # Need to use the correct data YAML (Dehazed)
        # We can reuse the one created for 'yolo_aod_advanced' if consistent
        # Or just pass the directory path and let helper create it
        res_ours = benchmark_model_custom(
            "data/RTTS_Dehazed_Advanced_Perceptual",
            "yolo_ours_finetuned",
            weights_ours
        )
        if res_ours:
            results["Ours"] = res_ours
    else:
        print("❌ Error: Could not find trained weights for Ours.")

    # --- Generate Markdown Tables ---
    
    # Table 1: Ablation / Algorithm Comparison
    # Rows: YOLO, YOLO+AOD, YOLO+Ours
    
    print("\n\n")
    print("### Table 1: Average Detection Precision (Ablation Study)")
    print("| Methods | Car | Bus | Motorcycle | mAP (3-Class) |")
    print("| :--- | :--- | :--- | :--- | :--- |")
    
    # We want to show:
    # 1. YOLO (Foggy Baseline)
    # 2. YOLO + AOD-Net (Original AOD)
    # 3. YOLO + AOD-Advanced (Ours Dehazed Only) -> Corresponds to "本文算法" in Fig 1 context of dehazing
    
    if "YOLO" in results:
        r = results["YOLO"]
        print(f"| YOLOv10 | {r['Car']:.3f} | {r['Bus']:.3f} | {r['Motorcycle']:.3f} | {r['mAP']:.3f} |")
        
    if "YOLO+AOD" in results:
        r = results["YOLO+AOD"]
        print(f"| YOLOv10 + AOD-Net | {r['Car']:.3f} | {r['Bus']:.3f} | {r['Motorcycle']:.3f} | {r['mAP']:.3f} |")
        
    if "YOLO+Advanced" in results:
        r = results["YOLO+Advanced"]
        print(f"| YOLOv10 + 本文算法 (Ours) | {r['Car']:.3f} | {r['Bus']:.3f} | {r['Motorcycle']:.3f} | {r['mAP']:.3f} |")
            
    # Table 2: Method Comparison (Before vs After)
    # Rows: YOLOv8, Improved YOLOv8
    
    print("\n\n")
    print("### Table 2: Average Detection Precision (Final Method Comparison)")
    print("| Methods | Car | Bus | Motorcycle | mAP (3-Class) |")
    print("| :--- | :--- | :--- | :--- | :--- |")
    
    # Row 1: YOLO (Baseline Foggy)
    if "YOLO" in results:
        r = results["YOLO"]
        print(f"| YOLOv10 | {r['Car']:.3f} | {r['Bus']:.3f} | {r['Motorcycle']:.3f} | {r['mAP']:.3f} |")
        
    # Row 2: Improved YOLO (Ours Finetuned)
    if "Ours" in results:
        r = results["Ours"]
        print(f"| 改进的 YOLOv10 (Ours) | {r['Car']:.3f} | {r['Bus']:.3f} | {r['Motorcycle']:.3f} | {r['mAP']:.3f} |")
        
    print("\nDone.")

if __name__ == "__main__":
    generate_final_tables()