实验3：Oriented R-CNN 旋转框检测

# 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

# dota_summary 解析 DOTA labelTxt；build_dota_teaching_subset 复制配对文件，
# 生成可追溯的小规模 trainval/val 目录，不会修改原始 DOTA。
from rsdet.notebook_utils import dota_summary
MMROTATE_ROOT = Path(os.getenv("MMROTATE_ROOT", str(ROOT / "third_party" / "mmrotate-v0.3.3")))
SOURCE_DOTA_ROOT = Path(os.getenv("DOTA_ROOT", str(ROOT / "data" / "raw" / "DOTA")))
SOURCE_TRAIN_ROOT = Path(os.getenv("DOTA_TRAIN_ROOT", str(SOURCE_DOTA_ROOT / "trainval")))
SOURCE_VAL_ROOT = Path(os.getenv("DOTA_VAL_ROOT", str(SOURCE_DOTA_ROOT / "val")))
SUBSET_DOTA_ROOT = RUN_ROOT / "oriented_rcnn_subset"
train_subset = build_dota_teaching_subset(
    SOURCE_TRAIN_ROOT / "images", SOURCE_TRAIN_ROOT / "labelTxt", SUBSET_DOTA_ROOT,
    max_images=SUBSET_MAX_IMAGES, seed=SEED, split="trainval",
)
val_subset = None
if (SOURCE_VAL_ROOT / "images").is_dir() and (SOURCE_VAL_ROOT / "labelTxt").is_dir():
    val_subset = build_dota_teaching_subset(
        SOURCE_VAL_ROOT / "images", SOURCE_VAL_ROOT / "labelTxt", SUBSET_DOTA_ROOT,
        max_images=max(1, SUBSET_MAX_IMAGES // 5), seed=SEED, split="val",
    )
DOTA_ROOT = SUBSET_DOTA_ROOT
ANN_DIR = DOTA_ROOT / "trainval" / "labelTxt"
os.environ["DOTA_ROOT"] = str(DOTA_ROOT)
os.environ["DOTA_TRAIN_ROOT"] = str(DOTA_ROOT / "trainval")
if val_subset:
    os.environ["DOTA_VAL_ROOT"] = str(DOTA_ROOT / "val")
ARTIFACTS["subset_manifest_trainval"] = str(train_subset["manifest"])
if val_subset:
    ARTIFACTS["subset_manifest_val"] = str(val_subset["manifest"])
CFG = ROOT / "configs" / "oriented_rcnn_remote_sensing.py"
print("完整 DOTA 已处理为教学子集:", {"trainval": train_subset, "val": val_subset}, {"mmrotate_root": str(MMROTATE_ROOT), "dota_root": str(DOTA_ROOT), "annotation_dir": str(ANN_DIR), "config": str(CFG)})

if ANN_DIR.exists():
    dota = dota_summary(ANN_DIR); print("DOTA 摘要:", dota)
    if dota["invalid_rows"]: raise ValueError(f"非法标注行: {dota['invalid_rows']}")
else:
    print("未发现真实 labelTxt；当前完成目录和配置审计。")

# 先检查每幅原图是否存在同名标注，防止训练时静默丢弃样本。
image_dir = DOTA_ROOT / "trainval" / "images"
image_files = [p for p in image_dir.glob("*") if p.suffix.lower() in {".jpg", ".png", ".tif", ".tiff"}] if image_dir.exists() else []
annotation_stems = {p.stem for p in ANN_DIR.glob("*.txt")} if ANN_DIR.exists() else set()
image_stems = {p.stem for p in image_files}
pair_summary = {"images": len(image_files), "annotations": len(annotation_stems), "missing_annotations": sorted(image_stems - annotation_stems)[:10], "orphan_annotations": sorted(annotation_stems - image_stems)[:10]}
print("图像—标注配对:", pair_summary)
print("切片索引原则：原图级划分后再切片，并记录 source_image、tile_id、offset 和边界裁剪信息。")

# 读取配置文本只做关键字段审计；真正的模型实例由 MMRotate 的 train.py 构建。
if not CFG.exists(): raise FileNotFoundError(CFG)
cfg_text = CFG.read_text(encoding="utf-8")
print("配置行数:", len(cfg_text.splitlines()), "DOTA 类别数预期: 15")
print("切片前置：先按原始大图划分 train/val，再切片；禁止同一原图切片跨集合。")

import importlib.util, subprocess
# importlib.util 检查依赖是否可导入；subprocess 调用 MMRotate 官方训练脚本。
required = ["mmcv", "mmrotate"]
availability = {name: importlib.util.find_spec(name) is not None for name in required}
TRAIN_SCRIPT, TEST_SCRIPT = MMROTATE_ROOT / "tools" / "train.py", MMROTATE_ROOT / "tools" / "test.py"
print("依赖状态:", availability, "train.py:", TRAIN_SCRIPT.exists(), "test.py:", TEST_SCRIPT.exists())

def run_mmrotate(stage="train"):
    """根据 stage 组装并执行 MMRotate 官方 train.py 或 test.py 命令。"""
    script = TRAIN_SCRIPT if stage == "train" else TEST_SCRIPT
    if not script.exists(): raise FileNotFoundError(script)
    work_dir = RUN_ROOT / "oriented_rcnn"
    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=MMROTATE_ROOT, check=True)

if RUN_TRAINING:
    if not availability["mmrotate"] or not DOTA_ROOT.exists(): raise RuntimeError("MMRotate 或 DOTA 数据未准备好")
    run_mmrotate("train")
else: print("RUN_TRAINING=0：完成配置审计，跳过训练。")

if RUN_TRAINING and TEST_SCRIPT.exists(): run_mmrotate("test")
print("正式评估应记录 OBB-mAP、逐类别 AP、角度误差、边界漏检和水平框基线差异。")

print("推理输出应保存预测四边形、原图叠加图和角度/长宽交换错误样本。")

log_candidates = [RUN_ROOT / "oriented_rcnn" / "log.json", RUN_ROOT / "oriented_rcnn" / "metrics.json"]
log_file = next((p for p in log_candidates if p.exists()), None)
if log_file:
    try:
        import matplotlib.pyplot as plt
        rows = [json.loads(line) for line in log_file.read_text(encoding="utf-8").splitlines() if line.strip()]
        epochs = [r.get("epoch") for r in rows if r.get("epoch") is not None]
        losses = [r.get("loss") for r in rows if r.get("loss") is not None]
        if epochs and losses:
            fig, ax = plt.subplots(figsize=(7, 4)); ax.plot(epochs, losses, marker="o")
            ax.set(xlabel="epoch", ylabel="loss", title="Oriented R-CNN loss"); ax.grid(alpha=.3)
            curve = RUN_ROOT / "oriented_rcnn_curves.png"; fig.tight_layout(); fig.savefig(curve, dpi=150); plt.close(fig)
            ARTIFACTS["curve"] = str(curve); print("曲线已保存:", curve)
        else:
            print("日志存在，但尚未提供可绘制的 epoch/loss 字段。")
    except (ImportError, json.JSONDecodeError) as exc:
        print("日志曲线绘制跳过:", exc)
else:
    print("尚无 MMRotate 日志；正式训练后将绘制 loss、学习率、逐类 AP、角度误差和尺度—Recall 曲线。")

print("若角度跳变或长宽交换，先检查顶点顺序、角度范围和标注转换，不要先更换模型。")
save_json("oriented_rcnn_audit.json", {"annotation_dir": str(ANN_DIR), "config": str(CFG), "dependencies": availability})
print_manifest(RUN_ROOT, **ARTIFACTS)