# ── 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)