Blame for scripts/4_train_cluster_classifier.py

"""
VC-ACT 阶段 4: 训练聚类分类器 (部署 Router)

训练一个 4-cam ResNet-18 分类器，输入 4 个相机的首帧，输出簇 ID。
部署时: 实时 4-cam 图像 → 此分类器 → 簇 ID → 对应簇的 ACT → 动作

输入:
  - results/cluster_labels_k10.npy 或 cluster_labels_auto.npy
  - results/episode_meta.json
  - bag_dir/ (原始 bag 文件)

输出:
  - models/exp_k10/cluster_classifier/best.pth
  - models/exp_auto/cluster_classifier/best.pth

用法:
  /venv/mult-act/bin/python3 4_train_cluster_classifier.py --mode k10
  /venv/mult-act/bin/python3 4_train_cluster_classifier.py --mode auto
"""

import argparse
import json
import random
from pathlib import Path

import cv2
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader, random_split
from torchvision import models, transforms
from PIL import Image

from rosbags.highlevel import AnyReader

CAM_TOPICS = [
    "/fisheye_left/image_raw/compressed",
    "/fisheye_right/image_raw/compressed",
    "/wide_mid/image_raw/compressed",
    "/wide_top/image_raw/compressed",
]
IMG_SIZE = (256, 256)


class FourCamClassifier(nn.Module):
    """4-camera ResNet-18 classifier.

    Each camera image goes through a shared ResNet-18 backbone (up to avgpool).
    4 x 512 features are concatenated and classified by a 2-layer MLP head.
    """

    def __init__(self, num_classes: int, pretrained: bool = True):
        super().__init__()
        weights = models.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None
        resnet = models.resnet18(weights=weights)
        self.backbone = nn.Sequential(*list(resnet.children())[:-1])  # -> (B, 512, 1, 1)
        self.head = nn.Sequential(
            nn.Linear(512 * 4, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, num_classes),
        )

    def forward(self, imgs: list[torch.Tensor]) -> torch.Tensor:
        """imgs: list of 4 tensors, each (B, 3, 224, 224)."""
        feats = []
        for img in imgs:
            f = self.backbone(img).squeeze(-1).squeeze(-1)  # (B, 512)
            feats.append(f)
        x = torch.cat(feats, dim=1)  # (B, 2048)
        return self.head(x)


def decode_compressed_image(msg) -> np.ndarray | None:
    """Decode a ROS CompressedImage to RGB (H, W, 3) uint8."""
    arr = np.frombuffer(msg.data, dtype=np.uint8)
    img_bgr = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if img_bgr is None:
        return None
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    return cv2.resize(img_rgb, IMG_SIZE, interpolation=cv2.INTER_LINEAR)


class BagFirstFrameDataset(Dataset):
    """Dataset: first frame of each bag from 4 cameras, with cluster label."""

    def __init__(self, bag_dir: Path, episode_meta: list, cluster_labels: np.ndarray,
                 transform=None):
        self.bag_dir = bag_dir
        self.meta = episode_meta
        self.labels = cluster_labels
        self.transform = transform

        # Pre-extract all first frames to memory (only 100 bags, 4 images each, ~100MB)
        print("Pre-loading first frames from bags...")
        self.frames = []  # list of {topic: RGB_array}
        for i, m in enumerate(self.meta):
            bag_path = self.bag_dir / m["bag_name"]
            cam_imgs = self._extract_first_frame(bag_path)
            self.frames.append(cam_imgs)
            if (i + 1) % 20 == 0:
                print(f"  [{i+1}/{len(self.meta)}]")
        print(f"Loaded {len(self.frames)} episodes")

    def _extract_first_frame(self, bag_path: Path) -> dict[str, np.ndarray]:
        """Extract first frame from each camera in a bag."""
        result = {}
        try:
            with AnyReader([bag_path]) as reader:
                cam_conns = [c for c in reader.connections if c.topic in CAM_TOPICS]
                for conn, _t, raw in reader.messages(connections=cam_conns):
                    if conn.topic not in result:
                        msg = reader.deserialize(raw, conn.msgtype)
                        img = decode_compressed_image(msg)
                        if img is not None:
                            result[conn.topic] = img
                    if len(result) == 4:
                        break
        except Exception as e:
            print(f"  Error reading {bag_path.name}: {e}")
        return result

    def __len__(self):
        return len(self.meta)

    def __getitem__(self, idx):
        cam_imgs = self.frames[idx]
        label = int(self.labels[idx])

        imgs = []
        for topic in CAM_TOPICS:
            if topic in cam_imgs:
                img_rgb = cam_imgs[topic]
            else:
                img_rgb = np.zeros((*IMG_SIZE, 3), dtype=np.uint8)

            if self.transform:
                img_pil = Image.fromarray(img_rgb)
                img_tensor = self.transform(img_pil)
            else:
                img_tensor = torch.from_numpy(img_rgb).permute(2, 0, 1).float() / 255.0
            imgs.append(img_tensor)

        return imgs, label


def collate_fn(batch):
    """Custom collate: list of (imgs_list, label) -> (4 batched tensors, labels)."""
    imgs_lists, labels = zip(*batch)
    batched_imgs = []
    for cam_idx in range(4):
        batched_imgs.append(torch.stack([imgs[cam_idx] for imgs in imgs_lists]))
    labels = torch.tensor(labels, dtype=torch.long)
    return batched_imgs, labels


def set_seed(seed=42):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def train(args):
    set_seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    vc_root = Path(args.vc_root)
    results_dir = vc_root / "results"
    bag_dir = Path(args.bag_dir)

    # Load labels
    if args.mode == "k10":
        labels = np.load(results_dir / "cluster_labels_k10.npy")
        save_dir = vc_root / "models" / "exp_k10" / "cluster_classifier"
    else:
        labels = np.load(results_dir / "cluster_labels_auto.npy")
        save_dir = vc_root / "models" / "exp_auto" / "cluster_classifier"

    with open(results_dir / "episode_meta.json") as f:
        meta = json.load(f)

    num_classes = int(labels.max()) + 1
    print(f"Mode: {args.mode}, num_classes={num_classes}, total_bags={len(labels)}")
    print(f"Label distribution: {np.bincount(labels).tolist()}")

    # Transforms
    train_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.RandomHorizontalFlip(),
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    val_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    # Load dataset (pre-loads all first frames)
    full_dataset = BagFirstFrameDataset(bag_dir, meta, labels, transform=train_transform)

    # Split
    total = len(full_dataset)
    val_size = max(1, int(total * args.val_ratio))
    train_size = total - val_size
    train_subset, val_subset = random_split(
        full_dataset, [train_size, val_size],
        generator=torch.Generator().manual_seed(args.seed)
    )

    # Override transform for val subset
    class TransformOverride(Dataset):
        def __init__(self, subset, transform):
            self.subset = subset
            self.transform = transform

        def __len__(self):
            return len(self.subset)

        def __getitem__(self, idx):
            # Get raw data from underlying dataset
            real_idx = self.subset.indices[idx]
            ds = self.subset.dataset
            cam_imgs = ds.frames[real_idx]
            label = int(ds.labels[real_idx])

            imgs = []
            for topic in CAM_TOPICS:
                if topic in cam_imgs:
                    img_rgb = cam_imgs[topic]
                else:
                    img_rgb = np.zeros((*IMG_SIZE, 3), dtype=np.uint8)
                img_pil = Image.fromarray(img_rgb)
                img_tensor = self.transform(img_pil)
                imgs.append(img_tensor)
            return imgs, label

    train_loader = DataLoader(
        train_subset, batch_size=args.batch_size, shuffle=True,
        num_workers=0, pin_memory=True, collate_fn=collate_fn
    )
    val_loader = DataLoader(
        TransformOverride(val_subset, val_transform),
        batch_size=args.batch_size, shuffle=False,
        num_workers=0, pin_memory=True, collate_fn=collate_fn
    )

    print(f"Train: {train_size}, Val: {val_size}")

    # Model
    model = FourCamClassifier(num_classes=num_classes, pretrained=True).to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)

    save_dir.mkdir(parents=True, exist_ok=True)
    best_val_acc = 0.0

    for epoch in range(1, args.epochs + 1):
        # Train
        model.train()
        train_loss, train_correct, train_total = 0.0, 0, 0
        for imgs, labels_batch in train_loader:
            imgs = [img.to(device) for img in imgs]
            labels_batch = labels_batch.to(device)

            optimizer.zero_grad()
            outputs = model(imgs)
            loss = criterion(outputs, labels_batch)
            loss.backward()
            optimizer.step()

            train_loss += loss.item() * labels_batch.size(0)
            train_correct += (outputs.argmax(1) == labels_batch).sum().item()
            train_total += labels_batch.size(0)

        scheduler.step()

        # Val
        model.eval()
        val_loss, val_correct, val_total = 0.0, 0, 0
        with torch.no_grad():
            for imgs, labels_batch in val_loader:
                imgs = [img.to(device) for img in imgs]
                labels_batch = labels_batch.to(device)
                outputs = model(imgs)
                loss = criterion(outputs, labels_batch)
                val_loss += loss.item() * labels_batch.size(0)
                val_correct += (outputs.argmax(1) == labels_batch).sum().item()
                val_total += labels_batch.size(0)

        train_acc = train_correct / max(train_total, 1)
        val_acc = val_correct / max(val_total, 1)

        print(f"Epoch {epoch:03d}/{args.epochs} | "
              f"Train Loss: {train_loss/max(train_total,1):.4f} Acc: {train_acc:.4f} | "
              f"Val Loss: {val_loss/max(val_total,1):.4f} Acc: {val_acc:.4f} | "
              f"LR: {scheduler.get_last_lr()[0]:.6f}")

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save({
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "val_acc": val_acc,
                "num_classes": num_classes,
                "mode": args.mode,
            }, save_dir / "best.pth")
            print(f"  -> Saved best (val_acc={val_acc:.4f})")

    # Save final
    torch.save({
        "epoch": args.epochs,
        "model_state_dict": model.state_dict(),
        "val_acc": val_acc,
        "num_classes": num_classes,
        "mode": args.mode,
    }, save_dir / "final.pth")

    print(f"\nDone. Best val acc: {best_val_acc:.4f}")
    print(f"Saved to: {save_dir}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="VC-ACT: Train 4-cam cluster classifier")
    parser.add_argument("--mode", type=str, choices=["k10", "auto"], required=True)
    parser.add_argument("--vc_root", type=str,
                        default="/workspace/Mult-skill ACT/VC-ACT")
    parser.add_argument("--bag_dir", type=str,
                        default="/workspace/Mult-skill ACT/data/26-03-04-pick_hanger")
    parser.add_argument("--epochs", type=int, default=50)
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--lr", type=float, default=1e-3)
    parser.add_argument("--val_ratio", type=float, default=0.2)
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()
    train(args)