| """Minimal inference for FairFace-MonkeyCake-3Task-224.""" |
| import torch |
| import timm |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from PIL import Image |
| from torchvision import transforms |
| from torchvision.transforms import InterpolationMode |
|
|
| ARCH = "eva02_large_patch14_224.mim_in22k" |
| FEATURE_DIM = 1024 |
| SIZE = 224 |
| MEAN = (0.48145466, 0.4578275, 0.40821073) |
| STD = (0.26862954, 0.26130258, 0.27577711) |
| RACE = ('White', 'Black', 'Latino_Hispanic', 'East Asian', 'Southeast Asian', 'Indian', 'Middle Eastern') |
| GENDER = ("Male", "Female") |
| AGE = ("0-2", "3-9", "10-19", "20-29", "30-39", "40-49", "50-59", "60-69", "more than 70") |
|
|
|
|
| class Student(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.backbone = timm.create_model(ARCH, pretrained=False, num_classes=0) |
| self.race_head = nn.Linear(FEATURE_DIM, 7) |
| self.gender_head = nn.Linear(FEATURE_DIM, 2) |
| self.age_head = nn.Linear(FEATURE_DIM, 8) |
|
|
| def forward(self, x): |
| f = self.backbone(x) |
| return {"race": self.race_head(f), "gender": self.gender_head(f), "age": self.age_head(f)} |
|
|
|
|
| def load(path="model.pt"): |
| ckpt = torch.load(path, map_location="cpu", weights_only=False) |
| model = Student() |
| model.load_state_dict(ckpt["ema"] if ckpt.get("selected_ema") else ckpt["model"]) |
| return model.eval() |
|
|
|
|
| def predict(model, image_path): |
| transform = transforms.Compose([ |
| transforms.Resize((SIZE, SIZE), interpolation=InterpolationMode.BICUBIC), |
| transforms.ToTensor(), |
| transforms.Normalize(mean=MEAN, std=STD), |
| ]) |
| with torch.no_grad(): |
| tensor = transform(Image.open(image_path).convert("RGB")).unsqueeze(0) |
| out = model(tensor) |
| age_index = int((out["age"][0] > 0).sum()) |
| return {"race": RACE[int(out["race"][0].argmax())], |
| "race_confidence": float(F.softmax(out["race"][0], dim=0).max()), |
| "gender": GENDER[int(out["gender"][0].argmax())], |
| "age": AGE[age_index]} |
|
|
