first

train_faster_rcnn.py

@@ -0,0 +1,291 @@
+"""用 YOLO 格式数据集训练 torchvision Faster R-CNN。"""
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+import cv2
+import numpy as np
+import torch
+import yaml
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+
+from detectors import build_faster_rcnn, resolve_device
+
+ROOT = Path(__file__).resolve().parent
+
+
+class YoloDetDataset(Dataset):
+    """读取 Ultralytics YOLO txt 标签（class cx cy w h，归一化）。"""
+
+    def __init__(self, img_dir: Path, label_dir: Path, imgsz: int = 640) -> None:
+        self.img_paths = sorted(
+            [p for p in img_dir.iterdir() if p.suffix.lower() in {".jpg", ".jpeg", ".png", ".bmp"}]
+        )
+        self.label_dir = label_dir
+        self.imgsz = int(imgsz)
+
+    def __len__(self) -> int:
+        return len(self.img_paths)
+
+    def __getitem__(self, idx: int):
+        path = self.img_paths[idx]
+        bgr = cv2.imread(str(path))
+        if bgr is None:
+            raise RuntimeError(f"无法读取图像: {path}")
+        h0, w0 = bgr.shape[:2]
+        scale = self.imgsz / max(h0, w0)
+        nh, nw = max(1, int(round(h0 * scale))), max(1, int(round(w0 * scale)))
+        bgr = cv2.resize(bgr, (nw, nh), interpolation=cv2.INTER_LINEAR)
+        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
+        image = torch.from_numpy(rgb).permute(2, 0, 1).float() / 255.0
+
+        label_path = self.label_dir / f"{path.stem}.txt"
+        boxes = []
+        labels = []
+        if label_path.exists():
+            for line in label_path.read_text(encoding="utf-8").splitlines():
+                parts = line.strip().split()
+                if len(parts) < 5:
+                    continue
+                cls_id = int(float(parts[0]))
+                cx, cy, bw, bh = map(float, parts[1:5])
+                x1 = (cx - bw / 2.0) * nw
+                y1 = (cy - bh / 2.0) * nh
+                x2 = (cx + bw / 2.0) * nw
+                y2 = (cy + bh / 2.0) * nh
+                x1 = float(np.clip(x1, 0, nw - 1))
+                y1 = float(np.clip(y1, 0, nh - 1))
+                x2 = float(np.clip(x2, 0, nw - 1))
+                y2 = float(np.clip(y2, 0, nh - 1))
+                if x2 <= x1 or y2 <= y1:
+                    continue
+                boxes.append([x1, y1, x2, y2])
+                # Faster R-CNN 前景从 1 开始
+                labels.append(cls_id + 1)
+
+        if boxes:
+            boxes_t = torch.tensor(boxes, dtype=torch.float32)
+            labels_t = torch.tensor(labels, dtype=torch.int64)
+        else:
+            boxes_t = torch.zeros((0, 4), dtype=torch.float32)
+            labels_t = torch.zeros((0,), dtype=torch.int64)
+
+        target = {
+            "boxes": boxes_t,
+            "labels": labels_t,
+            "image_id": torch.tensor([idx]),
+            "area": (boxes_t[:, 3] - boxes_t[:, 1]) * (boxes_t[:, 2] - boxes_t[:, 0])
+            if len(boxes_t)
+            else torch.zeros((0,), dtype=torch.float32),
+            "iscrowd": torch.zeros((len(boxes_t),), dtype=torch.int64),
+        }
+        return image, target
+
+
+def collate_fn(batch):
+    return tuple(zip(*batch))
+
+
+def load_data_yaml(path: Path) -> dict:
+    data = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
+    root = Path(data["path"])
+    if not root.is_absolute():
+        root = (path.parent / root).resolve()
+    return {
+        "root": root,
+        "train": root / data.get("train", "images/train"),
+        "val": root / data.get("val", "images/val"),
+        "names": data.get("names", {0: "object"}),
+    }
+
+
+def load_cfg(path: Path) -> dict:
+    cfg = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
+    return cfg
+
+
+@torch.inference_mode()
+def evaluate(model, loader, device) -> dict[str, float]:
+    model.eval()
+    total = 0
+    hit = 0
+    score_sum = 0.0
+    for images, targets in loader:
+        images = [img.to(device) for img in images]
+        outputs = model(images)
+        for out, tgt in zip(outputs, targets):
+            total += 1
+            gt = tgt["boxes"]
+            if len(gt) == 0:
+                continue
+            if len(out["boxes"]) == 0:
+                continue
+            # 简化指标：最高分框与任一 GT IoU>0.5 记命中
+            scores = out["scores"].detach().cpu()
+            boxes = out["boxes"].detach().cpu()
+            best = int(torch.argmax(scores))
+            score_sum += float(scores[best])
+            pb = boxes[best]
+            for gb in gt:
+                xx1 = max(float(pb[0]), float(gb[0]))
+                yy1 = max(float(pb[1]), float(gb[1]))
+                xx2 = min(float(pb[2]), float(gb[2]))
+                yy2 = min(float(pb[3]), float(gb[3]))
+                inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1)
+                area_p = max(0.0, float(pb[2] - pb[0])) * max(0.0, float(pb[3] - pb[1]))
+                area_g = max(0.0, float(gb[2] - gb[0])) * max(0.0, float(gb[3] - gb[1]))
+                union = area_p + area_g - inter + 1e-6
+                if inter / union >= 0.5:
+                    hit += 1
+                    break
+    return {
+        "val_images": float(total),
+        "recall_proxy": hit / max(1, total),
+        "avg_top_score": score_sum / max(1, total),
+    }
+
+
+def train_one_epoch(model, optimizer, loader, device, print_freq: int) -> float:
+    model.train()
+    loss_meter = 0.0
+    n = 0
+    pbar = tqdm(loader, desc="train", leave=False)
+    for i, (images, targets) in enumerate(pbar):
+        images = [img.to(device) for img in images]
+        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
+        loss_dict = model(images, targets)
+        losses = sum(loss for loss in loss_dict.values())
+        optimizer.zero_grad(set_to_none=True)
+        losses.backward()
+        optimizer.step()
+        loss_meter += float(losses.detach().cpu())
+        n += 1
+        if (i + 1) % print_freq == 0:
+            pbar.set_postfix(loss=f"{loss_meter / n:.4f}")
+    return loss_meter / max(1, n)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Train Faster R-CNN on YOLO-format dataset")
+    parser.add_argument("--cfg", type=str, default=str(ROOT / "faster_rcnn.yaml"))
+    args = parser.parse_args()
+
+    cfg = load_cfg(Path(args.cfg))
+    data_yaml = Path(cfg.get("data", "01data.yaml"))
+    if not data_yaml.is_absolute():
+        data_yaml = ROOT / data_yaml
+    data = load_data_yaml(data_yaml)
+
+    train_img = Path(data["train"])
+    val_img = Path(data["val"])
+    # YOLO layout: images/train -> labels/train
+    train_lbl = train_img.parent.parent / "labels" / train_img.name
+    val_lbl = val_img.parent.parent / "labels" / val_img.name
+    if not train_lbl.exists():
+        train_lbl = Path(str(train_img).replace("images", "labels"))
+    if not val_lbl.exists():
+        val_lbl = Path(str(val_img).replace("images", "labels"))
+
+    imgsz = int(cfg.get("imgsz", 640))
+    train_set = YoloDetDataset(train_img, train_lbl, imgsz=imgsz)
+    val_set = YoloDetDataset(val_img, val_lbl, imgsz=imgsz)
+    if len(train_set) == 0:
+        raise RuntimeError(f"训练集为空: {train_img}")
+
+    device = resolve_device(cfg.get("device", "0"))
+    num_classes = int(cfg.get("num_classes", 2))
+    model = build_faster_rcnn(num_classes, pretrained_backbone=bool(cfg.get("pretrained", True)))
+    model.to(device)
+
+    resume = str(cfg.get("resume") or "").strip()
+    start_epoch = 1
+    if resume:
+        ckpt = torch.load(resume, map_location="cpu", weights_only=False)
+        model.load_state_dict(ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt)
+        start_epoch = int(ckpt.get("epoch", 0)) + 1 if isinstance(ckpt, dict) else 1
+        print(f"已从 {resume} 恢复，下一轮 epoch={start_epoch}")
+
+    workers = int(cfg.get("num_workers", 2))
+    batch = int(cfg.get("batch_size", 2))
+    train_loader = DataLoader(
+        train_set,
+        batch_size=batch,
+        shuffle=True,
+        num_workers=workers,
+        collate_fn=collate_fn,
+        pin_memory=device.type == "cuda",
+    )
+    val_loader = DataLoader(
+        val_set,
+        batch_size=1,
+        shuffle=False,
+        num_workers=max(0, workers // 2),
+        collate_fn=collate_fn,
+    )
+
+    params = [p for p in model.parameters() if p.requires_grad]
+    optimizer = torch.optim.SGD(
+        params,
+        lr=float(cfg.get("lr", 0.005)),
+        momentum=float(cfg.get("momentum", 0.9)),
+        weight_decay=float(cfg.get("weight_decay", 0.0005)),
+    )
+    lr_scheduler = torch.optim.lr_scheduler.StepLR(
+        optimizer,
+        step_size=int(cfg.get("step_size", 12)),
+        gamma=float(cfg.get("gamma", 0.1)),
+    )
+
+    out_dir = ROOT / cfg.get("project", "runs/faster_rcnn") / cfg.get("name", "exp")
+    if out_dir.exists() and not bool(cfg.get("exist_ok", True)):
+        raise RuntimeError(f"输出目录已存在: {out_dir}")
+    out_dir.mkdir(parents=True, exist_ok=True)
+    weights_dir = out_dir / "weights"
+    weights_dir.mkdir(exist_ok=True)
+
+    epochs = int(cfg.get("epochs", 30))
+    print_freq = int(cfg.get("print_freq", 20))
+    eval_every = int(cfg.get("eval_every", 1))
+    best_score = -1.0
+
+    print(f"device={device}  train={len(train_set)}  val={len(val_set)}  out={out_dir}")
+    for epoch in range(start_epoch, epochs + 1):
+        t0 = time.perf_counter()
+        avg_loss = train_one_epoch(model, optimizer, train_loader, device, print_freq)
+        lr_scheduler.step()
+        metrics = {"loss": avg_loss}
+        if epoch % eval_every == 0 and len(val_set) > 0:
+            metrics.update(evaluate(model, val_loader, device))
+        elapsed = time.perf_counter() - t0
+        print(
+            f"epoch {epoch}/{epochs}  loss={avg_loss:.4f}  "
+            f"recall_proxy={metrics.get('recall_proxy', float('nan')):.3f}  "
+            f"time={elapsed:.1f}s"
+        )
+
+        ckpt = {
+            "epoch": epoch,
+            "model": model.state_dict(),
+            "num_classes": num_classes,
+            "names": data["names"],
+            "cfg": cfg,
+            "metrics": metrics,
+        }
+        torch.save(ckpt, weights_dir / "last.pth")
+        if epoch % 5 == 0:
+            torch.save(ckpt, weights_dir / f"epoch{epoch}.pth")
+        score = float(metrics.get("recall_proxy", 0.0))
+        if score >= best_score:
+            best_score = score
+            torch.save(ckpt, weights_dir / "best.pth")
+            print(f"  saved best.pth (recall_proxy={best_score:.3f})")
+
+    print(f"训练完成。权重目录: {weights_dir}")
+
+
+if __name__ == "__main__":
+    main()