CAGE_expression_inference-apvit/models/AffectNet8_Maxvit_VA/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


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)


# **** 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"
        )
        image = Image.open(image_path)

        classes = torch.tensor(self.dataframe.iloc[idx, 1], dtype=torch.int8)
        valence = torch.tensor(self.dataframe.iloc[idx, 2], dtype=torch.float16)
        arousal = torch.tensor(self.dataframe.iloc[idx, 3], dtype=torch.float16)

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

        return image, classes, valence, arousal

    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():
    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 APViTVAModel(nn.Module):
    def __init__(self):
        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.Dropout(0.3),
            nn.Linear(APVIT_EMBED_DIM, 2, bias=False),
        )

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


# Initialize the model
MODEL = APViTVAModel()
MODEL.to(DEVICE)

# **** Test the model performance for classification ****

# Set the model to evaluation mode
MODEL.load_state_dict(torch.load("model.pt", map_location=DEVICE))
MODEL.to(DEVICE)
MODEL.eval()

all_val_true_values = []
all_val_predicted_values = []
all_aro_true_values = []
all_aro_predicted_values = []

# Start inference on test set
with torch.no_grad():
    for images, _, val_true, aro_true in valid_loader:
        images, val_true, aro_true = (
            images.to(DEVICE),
            val_true.to(DEVICE),
            aro_true.to(DEVICE),
        )

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

        # Append to the lists --> Regression
        true_val_values = val_true.cpu().numpy()
        true_aro_values = aro_true.cpu().numpy()
        pred_val_values = val_pred.cpu().numpy()
        pred_aro_values = aro_pred.cpu().numpy()
        all_val_true_values.extend(true_val_values)
        all_aro_true_values.extend(true_aro_values)
        all_val_predicted_values.extend(pred_val_values)
        all_aro_predicted_values.extend(pred_aro_values)
df = pd.DataFrame(
    {
        "val_pred": all_val_predicted_values,
        "val_true": all_val_true_values,
        "aro_pred": all_aro_predicted_values,
        "aro_true": all_aro_true_values,
    }
)
df.to_csv("inference.csv", index=False)