orisys
292 lines
10 KiB
Python
292 lines
10 KiB
Python
"""用 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()
|