实验5：轻量化 YOLO 与 ONNX 推理

# 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

# 轻量化实验仍需使用真实图像和标签，只减少教学运行的图像对数量。
SOURCE_ROOT = Path(os.getenv("DATASET_ROOT", str(DATA_ROOT)))
SOURCE_IMAGE_DIR = Path(os.getenv("IMAGE_DIR", str(SOURCE_ROOT / "images")))
SOURCE_LABEL_DIR = Path(os.getenv("LABEL_DIR", str(SOURCE_ROOT / "labels")))
SOURCE_DATA_YAML = Path(os.getenv("DATA_YAML", str(SOURCE_ROOT / "data.yaml")))
subset_info = build_yolo_teaching_subset(
    SOURCE_IMAGE_DIR, SOURCE_LABEL_DIR, RUN_ROOT / "lightweight_subset",
    max_images=SUBSET_MAX_IMAGES, seed=SEED,
)
DATASET_ROOT = Path(subset_info["root"])
DATA_YAML = DATASET_ROOT / "data.yaml"
names = [x.strip() for x in os.getenv("YOLO_NAMES", "").split(",") if x.strip()] or read_yolo_names(SOURCE_DATA_YAML)
DATA_YAML.write_text(
    f"path: {DATASET_ROOT}\ntrain: images/train\nval: images/val\n"
    f"names: {names!r}\nnc: {len(names)}\n", encoding="utf-8"
)
ARTIFACTS["subset_manifest"] = str(subset_info["manifest"])
MODEL_VARIANTS = [x.strip() for x in os.getenv("MODEL_VARIANTS", "yolov8n.pt,yolov8s.pt").split(",") if x.strip()]
IMGSZ_LIST = [int(x) for x in os.getenv("IMGSZ_LIST", "320,640").split(",")]
print("完整数据已处理为轻量化教学子集:", subset_info, {"data": str(DATA_YAML), "models": MODEL_VARIANTS, "imgsz": IMGSZ_LIST})

# 统计标签文件和目标行数，只读检查，不修改数据集。
label_dir = DATASET_ROOT / "labels"
label_files = list(label_dir.rglob("*.txt")) if label_dir.exists() else []
print({"label_files": len(label_files), "objects": sum(len(x.read_text(encoding="utf-8").splitlines()) for x in label_files)})

dataset_root = DATASET_ROOT
image_files = [p for p in (dataset_root / "images").rglob("*") if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".tif", ".tiff"}] if (dataset_root / "images").exists() else []
label_files = list((dataset_root / "labels").rglob("*.txt")) if (dataset_root / "labels").exists() else []
image_stems = {p.relative_to(dataset_root / "images").with_suffix("") for p in image_files}
label_stems = {p.relative_to(dataset_root / "labels").with_suffix("") for p in label_files}
# 数据按 split 分目录时，仅保留 train/val/test 中与图像同层级的标签。
split_names = {"train", "val", "test", "trainval"}
image_splits = {key.parts[0] for key in image_stems if key.parts and key.parts[0] in split_names}
if image_splits:
    label_stems = {key for key in label_stems if key.parts and key.parts[0] in image_splits}
pair_summary = {"images": len(image_files), "labels": len(label_files), "missing_labels": sorted(image_stems - label_stems)[:10], "orphan_labels": sorted(label_stems - image_stems)[:10]}
print("轻量化实验数据配对:", pair_summary)
if pair_summary["missing_labels"] or pair_summary["orphan_labels"]:
    raise ValueError("发现图像与标签文件名不匹配，请先修复数据集目录。")

try:
    # YOLO 封装轻量模型的构建、训练、验证、推理和 ONNX 导出接口。
    from ultralytics import YOLO
    ULTRALYTICS_READY = True
except ImportError as exc:
    ULTRALYTICS_READY = False; print("Ultralytics 未安装:", exc)

def benchmark_model(model, image_path, imgsz, repeats=10):
    """预热后固定图像和后端，返回平均延迟（毫秒）与 FPS。"""
    for _ in range(2): model.predict(source=str(image_path), imgsz=imgsz, verbose=False)
    t0 = time.perf_counter()
    for _ in range(repeats): model.predict(source=str(image_path), imgsz=imgsz, verbose=False)
    latency = (time.perf_counter() - t0) / repeats
    return {"latency_ms": latency * 1000, "fps": 1 / max(latency, 1e-9)}
print("测速原则：warmup 后固定图像、输入尺寸、batch 和后端。")

benchmark_rows = []; models = {}
# 只从完整数据目录选取示例图像，不根据示例图推断总体性能。
first_image = find_first_image(DATASET_ROOT / "images")
if RUN_TRAINING and ULTRALYTICS_READY and first_image:
    for variant in MODEL_VARIANTS:
        model = YOLO(variant); models[variant] = model
        model.train(data=str(DATA_YAML), epochs=int(os.getenv("EPOCHS", "1")), imgsz=IMGSZ_LIST[0], batch=int(os.getenv("BATCH", "2")), workers=0, seed=SEED, project=str(RUN_ROOT), name=f"light_{Path(variant).stem}", exist_ok=True, plots=True)
else: print("RUN_TRAINING=0、依赖缺失或没有图像：跳过训练。")

for variant, model in models.items():
    validation = model.val(data=str(DATA_YAML), imgsz=IMGSZ_LIST[0], plots=True)
    print(variant, getattr(validation, "results_dict", {}))

for variant, model in models.items():
    if first_image: model.predict(source=str(first_image), save=True, project=str(RUN_ROOT), name=f"predict_{Path(variant).stem}", exist_ok=True)
    exported = model.export(format="onnx", simplify=False); print(variant, "ONNX:", exported)

if models and first_image:
    for variant, model in models.items():
        for imgsz in IMGSZ_LIST: benchmark_rows.append({"model": variant, "imgsz": imgsz, **benchmark_model(model, first_image, imgsz)})
save_json("lightweight_benchmark.json", benchmark_rows)
print("基准结果:", benchmark_rows)

print("报告需同时比较 mAP、Recall、参数量、模型文件大小、延迟、FPS、输入尺寸和部署后端。")
print_manifest(RUN_ROOT, **ARTIFACTS)