CAGE_expression_inference-apvit/models/AffectNet8_Maxvit_Combined/generate_csv.py

# ── APViT repo path setup (must precede mmcls/mmcv imports) ─────────────────
import sys
from pathlib import Path as _Path
_apvit_path = str(_Path(__file__).resolve().parent.parent.parent / "APViT")
if _apvit_path not in sys.path:
    sys.path.insert(0, _apvit_path)
# ─────────────────────────────────────────────────────────────────────────────

import pandas as pd
import os
import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
from PIL import Image
from pathlib import Path
import mmcv
from mmcls.models import build_classifier

SCRIPT_DIR = Path(__file__).resolve().parent
MODEL_PATH = os.getenv("MODEL_PATH", str(SCRIPT_DIR / "model.pt"))
APVIT_REPO = SCRIPT_DIR.parent.parent / "APViT"
APVIT_CONFIG = APVIT_REPO / "configs" / "apvit" / "RAF.py"


def resolve_images_dir(split_name, explicit_env_var, data_root, extract_root):
    explicit = os.getenv(explicit_env_var, "").strip()
    candidates = []
    if explicit:
        candidates.append(Path(explicit))

    candidates.append(Path(data_root) / f"{split_name}_set" / "images")
    candidates.append(Path(extract_root) / f"{split_name}_extracted" / "images")

    extract_split_root = Path(extract_root) / f"{split_name}_extracted"
    if extract_split_root.exists():
        for p in extract_split_root.rglob("images"):
            if p.is_dir():
                candidates.append(p)

    for p in candidates:
        if p.exists() and p.is_dir():
            return str(p)

    tried = "\n".join([str(p) for p in candidates])
    raise FileNotFoundError(
        f"Could not find images folder for split='{split_name}'. Tried:\n{tried}"
    )

# Load the annotations for validation from CSV file
DATA_ROOT = os.getenv("AFFECTNET_ROOT", "/workspace/data_affectnet/AffectNet")
EXTRACT_ROOT = os.getenv("AFFECTNET_EXTRACT_ROOT", f"{DATA_ROOT}/extracted")
ANNO_ROOT = os.getenv("AFFECTNET_ANNO_ROOT", "../../affectnet_annotations")
IMAGE_FOLDER = resolve_images_dir("train", "AFFECTNET_TRAIN_IMAGES", DATA_ROOT, EXTRACT_ROOT)
IMAGE_FOLDER_TEST = resolve_images_dir("val", "AFFECTNET_VAL_IMAGES", DATA_ROOT, EXTRACT_ROOT)
valid_annotations_path = os.getenv(
    "AFFECTNET_VAL_ANNO", f"{ANNO_ROOT}/val_set_annotation_without_lnd.csv"
)
valid_annotations_df = pd.read_csv(valid_annotations_path)

# Set parameters
BATCHSIZE = int(os.getenv("BATCHSIZE", "64"))
NUM_WORKERS = int(os.getenv("NUM_WORKERS", "0"))
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("DATA_ROOT =", DATA_ROOT)
print("EXTRACT_ROOT =", EXTRACT_ROOT)
print("AFFECTNET_TRAIN_IMAGES =", IMAGE_FOLDER)
print("AFFECTNET_VAL_IMAGES =", IMAGE_FOLDER_TEST)
print("MODEL_PATH =", MODEL_PATH)


# **** Create dataset and data loaders ****
class CustomDataset(Dataset):
    def __init__(self, dataframe, root_dir, transform=None, balance=False):
        self.dataframe = dataframe
        self.transform = transform
        self.root_dir = root_dir
        self.balance = balance

        if self.balance:
            self.dataframe = self.balance_dataset()

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

    def __getitem__(self, idx):
        image_path = os.path.join(
            self.root_dir, f"{self.dataframe['number'].iloc[idx]}.jpg"
        )
        if os.path.exists(image_path):
            image = Image.open(image_path)
        else:
            image = Image.new(
                "RGB", (112, 112), color="white"
            )  # Handle missing image file

        classes = torch.tensor(self.dataframe["exp"].iloc[idx], dtype=torch.long)
        labels = torch.tensor(self.dataframe.iloc[idx, 2:4].values, dtype=torch.float32)

        if self.transform:
            image = self.transform(image)

        return image, classes, labels

    def balance_dataset(self):
        balanced_df = self.dataframe.groupby("exp", group_keys=False).apply(
            lambda x: x.sample(self.dataframe["exp"].value_counts().min())
        )
        return balanced_df


transform_valid = transforms.Compose(
    [
        transforms.Resize(112),  # APViT / IR-50 requires 112x112 input
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)

valid_dataset = CustomDataset(
    dataframe=valid_annotations_df,
    root_dir=IMAGE_FOLDER_TEST,
    transform=transform_valid,
    balance=False,
)
valid_loader = DataLoader(
    valid_dataset, batch_size=BATCHSIZE, shuffle=False, num_workers=NUM_WORKERS
)

# ***** Define the model *****

APVIT_EMBED_DIM = 768


def _build_apvit_backbone():
    """Build APViT PoolingVitClassifier (backbone only, no head)."""
    cfg = mmcv.Config.fromfile(str(APVIT_CONFIG))
    cfg.model.pretrained = None
    cfg.model.head = None
    ir50_w = APVIT_REPO / "weights" / "backbone_ir50_ms1m_epoch120.pth"
    vit_small_w = APVIT_REPO / "weights" / "vit_small_p16_224-15ec54c9.pth"
    cfg.model.extractor.pretrained = str(ir50_w) if ir50_w.exists() else None
    cfg.model.vit.pretrained = str(vit_small_w) if vit_small_w.exists() else None
    return build_classifier(cfg.model)


class APViTWithHead(nn.Module):
    def __init__(self, num_outputs):
        super().__init__()
        self.apvit = _build_apvit_backbone()
        self.head = nn.Sequential(
            nn.LayerNorm(APVIT_EMBED_DIM),
            nn.Linear(APVIT_EMBED_DIM, APVIT_EMBED_DIM),
            nn.Tanh(),
            nn.Linear(APVIT_EMBED_DIM, num_outputs, bias=False),
        )

    def forward(self, x):
        features, _ = self.apvit.extract_feat(x)  # [B, 768] CLS token
        return self.head(features)


MODEL = APViTWithHead(num_outputs=10)
MODEL.to(DEVICE)  # Put the model to the GPU

# Set the model to evaluation mode
if not Path(MODEL_PATH).exists():
    raise FileNotFoundError(
        f"Missing Combined checkpoint: {MODEL_PATH}. "
        "Run train.py first or set MODEL_PATH to a valid checkpoint."
    )
MODEL.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
MODEL.to(DEVICE)
MODEL.eval()

all_labels_cls = []
all_predicted_cls = []

all_true_val = []
all_pred_val = []
all_true_aro = []
all_pred_aro = []

# Start inference on test set
with torch.no_grad():
    for images, classes, labels in iter(valid_loader):
        images, classes, labels = (
            images.to(DEVICE),
            classes.to(DEVICE),
            labels.to(DEVICE),
        )

        outputs = MODEL(images)
        outputs_cls = outputs[:, :8]
        outputs_reg = outputs[:, 8:]
        val_pred = outputs_reg[:, 0]
        aro_pred = outputs_reg[:, 1]

        _, predicted_cls = torch.max(outputs_cls, 1)

        all_labels_cls.extend(classes.cpu().numpy())
        all_predicted_cls.extend(predicted_cls.cpu().numpy())
        val_true = labels[:, 0]
        aro_true = labels[:, 1]

        all_true_val.extend(val_true.cpu().numpy())
        all_true_aro.extend(aro_true.cpu().numpy())
        all_pred_val.extend(val_pred.cpu().numpy())
        all_pred_aro.extend(aro_pred.cpu().numpy())

df = pd.DataFrame(
    {
        "cat_pred": all_predicted_cls,
        "cat_true": all_labels_cls,
        "val_pred": all_pred_val,
        "val_true": all_true_val,
        "aro_pred": all_pred_aro,
        "aro_true": all_true_aro,
    }
)
df.to_csv("inference.csv", index=False)