#!/usr/bin/env python3 # -*- coding: utf-8 -*- """训练 storefront / handover 二分类模型（MobileNetV3-Small）。""" import argparse import csv import os from pathlib import Path import torch import torch.nn as nn import yaml from PIL import Image from torch.utils.data import DataLoader, Dataset from torchvision import models, transforms from tqdm import tqdm class FrameDataset(Dataset): def __init__(self, rows, frames_dir, transform): self.rows = rows self.frames_dir = frames_dir self.transform = transform def __len__(self): return len(self.rows) def __getitem__(self, idx): path, label = self.rows[idx] img = Image.open(os.path.join(self.frames_dir, path)).convert("RGB") return self.transform(img), torch.tensor(label, dtype=torch.float32) def load_config(path): with open(path, encoding="utf-8") as f: cfg = yaml.safe_load(f) base = Path(path).resolve().parent for key in ("frames_dir", "labels_csv"): p = cfg["data"][key] if not os.path.isabs(p): cfg["data"][key] = str((base / p).resolve()) out = cfg["export"]["output_dir"] if not os.path.isabs(out): cfg["export"]["output_dir"] = str((base / out).resolve()) return cfg def read_labels(csv_path, task, split=None): """读取某任务的训练样本：该 task 行含 label 0/1；兼容旧版 task=other 负样本。""" rows = [] with open(csv_path, newline="", encoding="utf-8") as f: for r in csv.DictReader(f): if split and r.get("split") and r["split"] != split: continue row_task = r["task"] label = int(r["label"]) if row_task == task: rows.append((r["frame_path"], label)) elif row_task == "other" and label == 0: rows.append((r["frame_path"], 0)) return rows def count_labels(rows): pos = sum(1 for _, y in rows if y == 1) neg = len(rows) - pos return pos, neg def build_model(freeze_backbone=False): m = models.mobilenet_v3_small(weights=models.MobileNet_V3_Small_Weights.DEFAULT) if freeze_backbone: for p in m.features.parameters(): p.requires_grad = False in_f = m.classifier[0].in_features m.classifier = nn.Sequential( nn.Linear(in_f, 128), nn.Hardswish(), nn.Dropout(0.2), nn.Linear(128, 1), ) return m def build_transforms(size, augment=False): if augment: return transforms.Compose([ transforms.Resize((size, size)), transforms.RandomApply([ transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2), ], p=0.7), transforms.RandomApply([ transforms.GaussianBlur(kernel_size=3), ], p=0.2), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) return transforms.Compose([ transforms.Resize((size, size)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) def train_task(task, cfg): frames_dir = cfg["data"]["frames_dir"] train_rows = read_labels(cfg["data"]["labels_csv"], task, "train") val_rows = read_labels(cfg["data"]["labels_csv"], task, "val") if not train_rows: train_rows = read_labels(cfg["data"]["labels_csv"], task) if not val_rows: val_rows = train_rows[: max(1, len(train_rows) // 10)] train_pos, train_neg = count_labels(train_rows) val_pos, val_neg = count_labels(val_rows) print(f"[{task}] train pos={train_pos} neg={train_neg} | val pos={val_pos} neg={val_neg}") if train_neg == 0: print(f"WARN: [{task}] 训练集无负样本！请重新运行 prepare_dataset.py 生成 labels.csv") if train_pos < 10: print(f"WARN: [{task}] 正样本过少（<10），建议标注至少 80+ 条视频") size = cfg["model"]["image_size"] train_loader = DataLoader( FrameDataset(train_rows, frames_dir, build_transforms(size, augment=True)), batch_size=cfg["train"]["batch_size"], shuffle=True, num_workers=cfg["train"]["num_workers"], ) val_loader = DataLoader( FrameDataset(val_rows, frames_dir, build_transforms(size, augment=False)), batch_size=cfg["train"]["batch_size"], shuffle=False, num_workers=cfg["train"]["num_workers"], ) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") freeze = cfg["train"].get("freeze_backbone", True) and train_pos < 200 model = build_model(freeze_backbone=freeze).to(device) if freeze: print(f"[{task}] 小样本模式：冻结 backbone，仅训练分类头") pos_count = max(train_pos, 1) neg_count = max(train_neg, 1) auto_pos_weight = min(neg_count / pos_count, 10.0) pos_weight_val = cfg["model"].get("pos_weight") or auto_pos_weight pos_weight = torch.tensor([pos_weight_val], device=device) print(f"[{task}] pos_weight={pos_weight_val:.2f}") criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight) lr = cfg["train"]["lr"] if freeze: optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=lr) else: optimizer = torch.optim.Adam([ {"params": model.features.parameters(), "lr": lr * 0.1}, {"params": model.classifier.parameters(), "lr": lr}, ]) best_loss = float("inf") patience = 0 os.makedirs(cfg["export"]["output_dir"], exist_ok=True) ckpt_path = os.path.join(cfg["export"]["output_dir"], f"{task}.pt") for epoch in range(cfg["train"]["epochs"]): model.train() for x, y in tqdm(train_loader, desc=f"{task} epoch {epoch+1}"): x, y = x.to(device), y.to(device).unsqueeze(1) optimizer.zero_grad() loss = criterion(model(x), y) loss.backward() optimizer.step() model.eval() val_loss = 0.0 n = 0 with torch.no_grad(): for x, y in val_loader: x, y = x.to(device), y.to(device).unsqueeze(1) val_loss += criterion(model(x), y).item() * x.size(0) n += x.size(0) val_loss /= max(n, 1) print(f"{task} epoch {epoch+1} val_loss={val_loss:.4f}") if val_loss < best_loss: best_loss = val_loss patience = 0 torch.save(model.state_dict(), ckpt_path) else: patience += 1 if patience >= cfg["train"]["early_stop_patience"]: break print(f"保存 {ckpt_path}") return ckpt_path def main(): parser = argparse.ArgumentParser() parser.add_argument("-c", "--config", default=str(Path(__file__).parent / "config.yaml")) parser.add_argument("--task", choices=["storefront", "handover", "both"], default="both") args = parser.parse_args() cfg = load_config(args.config) tasks = ["storefront", "handover"] if args.task == "both" else [args.task] for t in tasks: train_task(t, cfg) if __name__ == "__main__": main()