实验1：Faster R-CNN 水平框检测

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, 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")))
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):
    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"):
    """Save loss/learning-rate curves when a real run produced history."""
    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

print("建议环境：Python 3.9、PyTorch 2.0、torchvision 0.15；GPU 训练优先使用 CUDA。")
print("本次运行目录:", RUN_ROOT.resolve())

# SOURCE_COCO_JSON 是 DIOR 转换得到的完整 COCO 标注；函数只读取源数据，
# 再把图像—标注配对复制到 RUN_ROOT 下的教学子集目录。
SOURCE_COCO_JSON = Path(os.getenv("COCO_JSON", str(DATA_ROOT / "annotations.json")))
SOURCE_IMAGE_DIR = Path(os.getenv("IMAGE_DIR", str(DATA_ROOT / "images")))
if not SOURCE_COCO_JSON.exists():
    candidates = [DATA_ROOT / "annotations" / "train.json", DATA_ROOT / "train.json"]
    SOURCE_COCO_JSON = next((path for path in candidates if path.exists()), SOURCE_COCO_JSON)
if not SOURCE_COCO_JSON.exists():
    raise FileNotFoundError(f"DIOR 的 COCO 标注不存在: {SOURCE_COCO_JSON.resolve()}")
if not SOURCE_IMAGE_DIR.exists():
    raise FileNotFoundError(f"DIOR 图像目录不存在: {SOURCE_IMAGE_DIR.resolve()}")
SUBSET_ROOT = RUN_ROOT / "faster_rcnn_subset"
subset_info = build_coco_teaching_subset(
    SOURCE_COCO_JSON, SUBSET_ROOT / "annotations.json",
    max_images=SUBSET_MAX_IMAGES, max_classes=SUBSET_MAX_CLASSES, seed=SEED,
    source_image_dir=SOURCE_IMAGE_DIR,
)
COCO_JSON = Path(subset_info["annotation"])
IMAGE_DIR = SUBSET_ROOT
ARTIFACTS["subset_manifest"] = str(subset_info["manifest"])
summary = coco_summary(COCO_JSON)
ARTIFACTS["coco_summary"] = summary
print("DIOR → COCO 教学子集准备完成:", subset_info, summary)

with COCO_JSON.open(encoding="utf-8") as stream:
    coco = json.load(stream)
print("图像数:", len(coco["images"]), "标注数:", len(coco["annotations"]), "类别:", coco["categories"])
ann_counts = {}
for ann in coco["annotations"]:
    ann_counts[ann["category_id"]] = ann_counts.get(ann["category_id"], 0) + 1
print("类别实例统计:", ann_counts)
first = find_first_image(IMAGE_DIR)
if first:
    try:
        from PIL import Image
        import matplotlib.pyplot as plt
        image = Image.open(first).convert("RGB")
        plt.figure(figsize=(5, 5)); plt.imshow(image); plt.title(first.name); plt.axis("off"); plt.show()
    except ImportError:
        print("Pillow 或 matplotlib 未安装，跳过图像显示。")

import numpy as np
from PIL import Image
try:
    import torch
    from torch.utils.data import Dataset, DataLoader
    TORCH_READY = True
except ImportError as exc:
    torch = None; Dataset = object; DataLoader = None; TORCH_READY = False
    print("PyTorch 未安装:", exc)

cat_ids = sorted(item["id"] for item in coco["categories"])
cat_to_label = {cat_id: i + 1 for i, cat_id in enumerate(cat_ids)}
images = {item["id"]: item for item in coco["images"]}
ann_by_image = {image_id: [] for image_id in images}
for ann in coco["annotations"]: ann_by_image.setdefault(ann["image_id"], []).append(ann)

class CocoDetection(Dataset):
    """把 COCO xywh 转为 torchvision xyxy；不在这里做随机增强，便于复核。"""
    def __init__(self, image_dir, records, annotations): self.image_dir, self.records, self.annotations = Path(image_dir), records, annotations
    def __len__(self): return len(self.records)
    def __getitem__(self, index):
        record = self.records[index]
        image = torch.as_tensor(np.asarray(Image.open(self.image_dir / record["file_name"]).convert("RGB")).transpose(2, 0, 1), dtype=torch.float32) / 255.0
        boxes, labels, areas = [], [], []
        for ann in self.annotations.get(record["id"], []):
            x, y, w, h = ann["bbox"]; boxes.append([x, y, x + w, y + h]); labels.append(cat_to_label[ann["category_id"]]); areas.append(w * h)
        target = {"boxes": torch.tensor(boxes, dtype=torch.float32), "labels": torch.tensor(labels, dtype=torch.int64), "image_id": torch.tensor([record["id"]]), "area": torch.tensor(areas, dtype=torch.float32), "iscrowd": torch.zeros(len(boxes), dtype=torch.int64)}
        return image, target

def collate_fn(batch): return tuple(zip(*batch))
records = list(images.values()); split = max(1, int(len(records) * .8))
if TORCH_READY:
    train_ds = CocoDetection(IMAGE_DIR, records[:split], ann_by_image); val_ds = CocoDetection(IMAGE_DIR, records[split:] or records[:1], ann_by_image)
    train_loader = DataLoader(train_ds, batch_size=int(os.getenv("BATCH", "2")), shuffle=True, collate_fn=collate_fn)
    val_loader = DataLoader(val_ds, batch_size=1, shuffle=False, collate_fn=collate_fn)
    print("训练/验证图像数:", len(train_ds), len(val_ds))
else:
    train_ds = val_ds = train_loader = val_loader = None

model = None; device = None
if TORCH_READY:
    try:
        from torchvision.models.detection import fasterrcnn_resnet50_fpn
        from torchvision.models.detection import FasterRCNN_ResNet50_FPN_Weights
        from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
        weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT if os.getenv("USE_PRETRAINED", "0") == "1" else None
        model = fasterrcnn_resnet50_fpn(weights=weights, min_size=320, max_size=640)
        model.roi_heads.box_predictor = FastRCNNPredictor(model.roi_heads.box_predictor.cls_score.in_features, len(cat_ids) + 1)
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu"); model.to(device)
        print("模型构建完成:", type(model).__name__, "device=", device)
    except ImportError as exc:
        print("torchvision detection 未就绪:", exc)

optimizer = scheduler = None
if model is not None:
    optimizer = torch.optim.SGD([p for p in model.parameters() if p.requires_grad], lr=float(os.getenv("LR", "0.005")), momentum=.9, weight_decay=.0005)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=.1)
    print("优化器:", type(optimizer).__name__, "初始学习率:", optimizer.param_groups[0]["lr"])

def train_one_epoch(model, loader, optimizer, device):
    model.train(); losses = []
    for image_batch, target_batch in loader:
        image_batch = [x.to(device) for x in image_batch]; target_batch = [{k: v.to(device) for k, v in t.items()} for t in target_batch]
        loss_dict = model(image_batch, target_batch); loss = sum(loss_dict.values())
        optimizer.zero_grad(); loss.backward(); optimizer.step(); losses.append(float(loss.detach().cpu()))
    return sum(losses) / max(1, len(losses))

def evaluate_predictions(model, loader, device):
    model.eval(); count = 0; boxes = 0
    with torch.no_grad():
        for image_batch, _ in loader:
            output = model([x.to(device) for x in image_batch]); count += len(output); boxes += sum(len(item["boxes"]) for item in output)
    return {"images": count, "predictions": boxes}

history = []
if model is None:
    print("模型不可用，完成数据链路检查后跳过训练。")
elif RUN_TRAINING:
    for epoch in range(int(os.getenv("EPOCHS", "1"))):
        loss = train_one_epoch(model, train_loader, optimizer, device); scheduler.step()
        row = {"epoch": epoch + 1, "loss": loss, "lr": optimizer.param_groups[0]["lr"]}; history.append(row); print(row)
    checkpoint = RUN_ROOT / "faster_rcnn_checkpoint.pt"
    torch.save({"model": model.state_dict(), "categories": summary["category_names"], "seed": SEED}, checkpoint)
    ARTIFACTS["checkpoint"] = str(checkpoint)
else:
    print("RUN_TRAINING=0：跳过长时间训练。")
save_json("faster_rcnn_history.json", history)

evaluation = {}
if model is not None:
    evaluation = evaluate_predictions(model, val_loader, device)
    print("教学子集验证集预测统计:", evaluation)
else:
    print("无模型输出；正式运行时应接入 COCOeval，报告 mAP50、mAP50:95、Precision、Recall。")
save_json("faster_rcnn_evaluation.json", evaluation)

inference_record = {"status": "skipped", "reason": "模型未训练或依赖未安装"}
if model is not None and RUN_TRAINING and first:
    model.eval()
    with torch.no_grad():
        image = torch.as_tensor(np.asarray(Image.open(first).convert("RGB")).transpose(2, 0, 1), dtype=torch.float32).to(device) / 255.0
        output = model([image])[0]
    inference_record = {"image": str(first), "boxes": int(len(output["boxes"])), "score_threshold": .25}
    print("推理结果:", inference_record)
save_json("faster_rcnn_inference.json", inference_record)

plot_history(history, "faster_rcnn_curves.png")
print("最终产物:", ARTIFACTS)
print_manifest(RUN_ROOT, **ARTIFACTS)