"""Standalone inference for the four strategies in the DFER thesis. The published checkpoints retain their original parameter names. No private repository, training code, dataset or ImageNet weight download is required. """ import argparse import re import types import timm import torch from huggingface_hub import HfApi, hf_hub_download from torch import nn from torchvision import transforms as T from torchvision.transforms import InterpolationMode REPO_ID = "headless-start/parameter-efficient-dfer" MODEL_NAME = "vit_base_patch16_224.augreg_in21k_ft_in1k" CLASS_NAMES = ["angry", "disgust", "fear", "happy", "sad", "surprise"] METHODS = ("full_ft", "linear_probe", "adaptformer", "lora") def build_transform(): """The evaluation transform; input must already be a PIL face crop.""" return T.Compose([ T.Grayscale(num_output_channels=3), T.Resize((224, 224), interpolation=InterpolationMode.BICUBIC), T.ToTensor(), T.Normalize((0.5,) * 3, (0.5,) * 3), ]) class _LoRAQKV(nn.Module): def __init__(self, qkv): super().__init__() self.qkv = qkv width, rank = qkv.in_features, 4 self.lora_q_A = nn.Linear(width, rank, bias=False) self.lora_q_B = nn.Linear(rank, width, bias=False) self.lora_v_A = nn.Linear(width, rank, bias=False) self.lora_v_B = nn.Linear(rank, width, bias=False) def forward(self, x): q, k, v = self.qkv(x).chunk(3, dim=-1) q = q + 2.0 * self.lora_q_B(self.lora_q_A(x)) v = v + 2.0 * self.lora_v_B(self.lora_v_A(x)) return torch.cat((q, k, v), dim=-1) class _AdaptMLP(nn.Module): def __init__(self, width): super().__init__() self.down = nn.Linear(width, 64) self.up = nn.Linear(64, width) self.drop = nn.Dropout(0.1) def forward(self, x): return 0.1 * self.up(self.drop(torch.relu(self.down(x)))) def _adapter_forward(block, x, attn_mask=None, is_causal=False): x = x + block.drop_path1(block.ls1( block.attn(block.norm1(x), attn_mask=attn_mask, is_causal=is_causal))) return (x + block.drop_path2(block.ls2(block.mlp(block.norm2(x)))) + block.adaptmlp(x)) def _build_model(): return timm.create_model( MODEL_NAME, pretrained=False, num_classes=6, drop_rate=0.0, drop_path_rate=0.1, ) def _resolve_revision(revision): # Resolve a branch once so backbone and fold come from the same commit. if re.fullmatch(r"[0-9a-f]{40}", revision): return revision return HfApi(token=False).model_info(REPO_ID, revision=revision).sha def _checkpoint(filename, revision, method, stage, fold=None): path = hf_hub_download(REPO_ID, filename, revision=revision, token=False) checkpoint = torch.load(path, map_location="cpu", weights_only=True) if (checkpoint.get("schema_version") != 3 or checkpoint.get("class_names") != CLASS_NAMES or checkpoint.get("method") != method or checkpoint.get("cfg", {}).get("data", {}).get("name") != stage or checkpoint.get("trainable_only") != (method != "full_ft")): raise ValueError(f"{filename}: incompatible checkpoint metadata") if fold is not None: identity = checkpoint.get("identity", {}) if (identity.get("unit_key") != f"fold_{fold:02d}" or identity.get("source_stage") != "ferplus"): raise ValueError(f"{filename}: wrong fold or source stage") state = checkpoint.get("model") if not isinstance(state, dict) or not state: raise ValueError(f"{filename}: missing model state") return state def load_stage1(*, device="cpu", revision="main"): """Load the complete selected six-class FER+ model in evaluation mode.""" revision = _resolve_revision(revision) state = _checkpoint("stage1/best.ckpt", revision, "full_ft", "ferplus") model = _build_model() model.load_state_dict(state, strict=True) return model.to(device).eval() def load_stage2(method, fold, *, device="cpu", revision="main"): """Load one selected KMU-FED fold (0–9), including its frozen encoder.""" if method not in METHODS: raise ValueError(f"method must be one of {METHODS}") if isinstance(fold, bool) or not isinstance(fold, int) or not 0 <= fold <= 9: raise ValueError("fold must be an integer from 0 to 9") revision = _resolve_revision(revision) filename = f"stage2/{method}/fold_{fold:02d}/best.ckpt" state = _checkpoint(filename, revision, method, "kmufed", fold) if method == "full_ft": model = _build_model() model.load_state_dict(state, strict=True) else: model = load_stage1(revision=revision) for parameter in model.parameters(): parameter.requires_grad = False if method == "lora": for block in model.blocks: block.attn.qkv = _LoRAQKV(block.attn.qkv) elif method == "adaptformer": for block in model.blocks: block.adaptmlp = _AdaptMLP(model.embed_dim) block.forward = types.MethodType(_adapter_forward, block) for parameter in model.head.parameters(): parameter.requires_grad = True expected = {name for name, p in model.named_parameters() if p.requires_grad} if set(state) != expected: raise ValueError(f"{filename}: incomplete or unexpected trained parameters") missing, unexpected = model.load_state_dict(state, strict=False) if unexpected or set(missing) != set(model.state_dict()) - expected: raise ValueError(f"{filename}: incompatible model state") return model.to(device).eval() def main(): from PIL import Image parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("image", help="path to an already cropped face image") parser.add_argument("--method", choices=METHODS, default="lora") parser.add_argument("--fold", type=int, choices=range(10), default=0) parser.add_argument("--device", default="cpu") parser.add_argument("--revision", default="main") args = parser.parse_args() model = load_stage2(args.method, args.fold, device=args.device, revision=args.revision) with Image.open(args.image) as image: x = build_transform()(image.convert("RGB")).unsqueeze(0).to(args.device) with torch.inference_mode(): probabilities = model(x).softmax(dim=1)[0].cpu() for index in probabilities.argsort(descending=True).tolist(): print(f"{CLASS_NAMES[index]}: {probabilities[index].item():.4f}") if __name__ == "__main__": main()