实验4：Deformable DETR 多尺度检测

# pathlib：跨平台处理项目路径；json：保存实验清单；os：读取环境变量；
# sys：把项目源码加入导入路径；time：测速；platform：记录运行平台。
from pathlib import Path
import json, os, sys, time, platform

ROOT = Path.cwd()
while not (ROOT / "pyproject.toml").exists() and ROOT != ROOT.parent:
    ROOT = ROOT.parent
sys.path.insert(0, str(ROOT / "src"))
from rsdet.notebook_utils import (
    set_seed,
    build_coco_teaching_subset,
    build_yolo_teaching_subset,
    build_dota_teaching_subset,
    read_yolo_names,
    coco_summary,
    print_manifest,
    find_first_image,
)

SEED = int(os.getenv("SEED", "42"))
set_seed(SEED)
RUN_TRAINING = bool(int(os.getenv("RUN_TRAINING", "0")))
# Notebook 读取完整数据源，再写出可追溯的教学子集；原始目录不被覆盖。
SUBSET_MAX_IMAGES = int(os.getenv("SUBSET_MAX_IMAGES", "3000"))
SUBSET_MAX_CLASSES = int(os.getenv("SUBSET_MAX_CLASSES", "6"))
DATA_ROOT = Path(os.getenv("DATA_ROOT", str(ROOT / "data" / "processed" / "dior")))
RUN_ROOT = ROOT / "runs" / "notebook_outputs"
RUN_ROOT.mkdir(parents=True, exist_ok=True)
ARTIFACTS = {"seed": SEED, "platform": platform.platform(), "python": platform.python_version()}
print({"root": str(ROOT), "run_training": RUN_TRAINING, "subset_max_images": SUBSET_MAX_IMAGES, "subset_max_classes": SUBSET_MAX_CLASSES, "seed": SEED})

def save_json(name, payload):
    """将阶段性结果写入 runs/notebook_outputs，方便报告引用和复核。"""
    path = RUN_ROOT / name
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2, default=str), encoding="utf-8")
    ARTIFACTS[name] = str(path)
    return path

def plot_history(history, name="training_curves.png"):
    """读取训练历史并绘制 loss、学习率曲线；无历史时只给出提示。"""
    if not history:
        print("暂无训练历史；设置 RUN_TRAINING=1 后会保存曲线。")
        return None
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        print("matplotlib 未安装，跳过曲线绘制。")
        return None
    epochs = [row.get("epoch", i + 1) for i, row in enumerate(history)]
    fig, axes = plt.subplots(1, 2, figsize=(10, 4))
    axes[0].plot(epochs, [row.get("loss", float("nan")) for row in history], marker="o")
    axes[0].set_title("loss"); axes[0].set_xlabel("epoch"); axes[0].grid(alpha=.3)
    axes[1].plot(epochs, [row.get("lr", float("nan")) for row in history], marker="o", color="darkorange")
    axes[1].set_title("learning rate"); axes[1].set_xlabel("epoch"); axes[1].grid(alpha=.3)
    fig.tight_layout()
    path = RUN_ROOT / name; fig.savefig(path, dpi=150); plt.close(fig)
    ARTIFACTS[name] = str(path)
    print("曲线已保存:", path)
    return path

# importlib.util 检查 MMDetection 依赖；subprocess 调用官方训练和测试脚本。
import importlib.util, subprocess
MMDET_ROOT = Path(os.getenv("MMDET_ROOT", str(ROOT / "third_party" / "mmdetection-v3.1.0")))
SOURCE_COCO_ROOT = Path(os.getenv("COCO_ROOT", str(ROOT / "data" / "processed" / "dior")))
COCO_ROOT = RUN_ROOT / "deformable_detr_subset"
COCO_ROOT.mkdir(parents=True, exist_ok=True)
source_train = SOURCE_COCO_ROOT / "annotations" / "train.json"
source_val = SOURCE_COCO_ROOT / "annotations" / "val.json"
train_subset = build_coco_teaching_subset(
    source_train, COCO_ROOT / "annotations" / "train.json",
    max_images=SUBSET_MAX_IMAGES, max_classes=SUBSET_MAX_CLASSES, seed=SEED,
    source_image_dir=SOURCE_COCO_ROOT,
    class_names=["airplane", "ship", "vehicle", "bridge", "storage-tank", "harbor"],
)
val_subset = build_coco_teaching_subset(
    source_val, COCO_ROOT / "annotations" / "val.json",
    max_images=max(1, SUBSET_MAX_IMAGES // 5), max_classes=SUBSET_MAX_CLASSES, seed=SEED,
    source_image_dir=SOURCE_COCO_ROOT,
    class_names=["airplane", "ship", "vehicle", "bridge", "storage-tank", "harbor"],
)
ARTIFACTS["subset_manifest_train"] = str(train_subset["manifest"])
ARTIFACTS["subset_manifest_val"] = str(val_subset["manifest"])
CFG = ROOT / "configs" / "deformable_detr_remote_sensing.py"
TRAIN_SCRIPT, TEST_SCRIPT = MMDET_ROOT / "tools" / "train.py", MMDET_ROOT / "tools" / "test.py"
print("完整 COCO 数据已处理为教学子集:", {"train": train_subset, "val": val_subset}, {"mmdetection_root": str(MMDET_ROOT), "data_root": str(COCO_ROOT), "config": str(CFG)})

# 先检查目录规模和文件组织，不改变完整数据集内容。
print("COCO 目录应包含 train/、val/ 和 annotations/train.json、annotations/val.json。")
if COCO_ROOT.exists(): print("COCO_ROOT 已存在，文件数:", sum(1 for _ in COCO_ROOT.rglob("*")))
else: print("未发现真实 COCO_ROOT；只执行配置审计。")

# COCO 摘要函数只读取元信息：图像数、标注数、类别数和类别名称。
annotation_candidates = [COCO_ROOT / "annotations" / "train.json", COCO_ROOT / "annotations" / "val.json", COCO_ROOT / "annotations.json"]
annotation_files = [p for p in annotation_candidates if p.exists()]
dataset_summary = {}
for annotation_file in annotation_files:
    try:
        dataset_summary[str(annotation_file)] = coco_summary(annotation_file)
    except (ValueError, KeyError, json.JSONDecodeError) as exc:
        dataset_summary[str(annotation_file)] = {"error": str(exc)}
print("COCO 数据集摘要:", dataset_summary or "尚无 COCO 标注文件")

# 配置审计用于提前发现类别数、数据根目录和训练轮数缺失。
cfg_text = CFG.read_text(encoding="utf-8") if CFG.exists() else ""
required = ["_base_", "data_root", "num_classes", "max_epochs"]
print("配置检查:", {token: token in cfg_text for token in required})

print("MMDetection 依赖状态:", {name: importlib.util.find_spec(name) is not None for name in ["mmengine", "mmdet"]})
def run_mmdet(stage="train"):
    """组装 MMDetection 官方 train.py/test.py 命令并返回子进程结果。"""
    script = TRAIN_SCRIPT if stage == "train" else TEST_SCRIPT
    if not script.exists(): raise FileNotFoundError(script)
    work_dir = RUN_ROOT / "deformable_detr"; cmd = [sys.executable, str(script), str(CFG)]
    if stage == "train": cmd += ["--work-dir", str(work_dir)]
    else: cmd += [str(work_dir / "latest.pth")]
    print("执行命令:", " ".join(map(str, cmd))); return subprocess.run(cmd, cwd=MMDET_ROOT, check=True)

if RUN_TRAINING:
    if not TRAIN_SCRIPT.exists() or not COCO_ROOT.exists(): raise RuntimeError("MMDetection、配置或 COCO 数据未准备好")
    run_mmdet("train")
else: print("RUN_TRAINING=0：跳过外部训练命令。")

if RUN_TRAINING and TEST_SCRIPT.exists(): run_mmdet("test")
print("正式报告应记录 COCO mAP、AP_small、收敛 epoch、FPS、显存和与 YOLOv8 的对照。")

print("推理结果应保存预测 JSON、带框图像和小目标错误样本。")

log_file = RUN_ROOT / "deformable_detr" / "vis_data" / "scalars.json"
if log_file.exists():
    try:
        import matplotlib.pyplot as plt
        rows = [json.loads(line) for line in log_file.read_text(encoding="utf-8").splitlines() if line.strip()]
        steps = [r.get("step") for r in rows if r.get("step") is not None]
        losses = [r.get("loss") for r in rows if r.get("loss") is not None]
        if steps and losses:
            fig, ax = plt.subplots(figsize=(7, 4)); ax.plot(steps, losses, linewidth=1.5)
            ax.set(xlabel="step", ylabel="loss", title="Deformable DETR loss"); ax.grid(alpha=.3)
            curve = RUN_ROOT / "deformable_detr_curves.png"; fig.tight_layout(); fig.savefig(curve, dpi=150); plt.close(fig)
            ARTIFACTS["curve"] = str(curve); print("曲线已保存:", curve)
        else:
            print("日志存在，但尚未提供可绘制的 step/loss 字段。")
    except (ImportError, json.JSONDecodeError) as exc:
        print("日志曲线绘制跳过:", exc)
else:
    print("尚无 MMDetection 日志；正式训练后读取 scalars.json 绘制 loss、学习率和 mAP 曲线。")

save_json("deformable_detr_audit.json", {"config": str(CFG), "data_root": str(COCO_ROOT), "mmdetection_root": str(MMDET_ROOT)})
print("优先排查 _base_ 路径、num_classes、ann_file、评估器路径和显存。")
print_manifest(RUN_ROOT, **ARTIFACTS)