👥 Pedestrian Age + Gender — v20260507 訓練中期報告

import torch, torch.nn as nn, timm
from PIL import Image
import torchvision.transforms as T

# Model 結構（從 train script 抽出）
class MultiHead(nn.Module):
    def __init__(self, backbone_name, drop_rate=0.3, num_age=4):
        super().__init__()
        self.backbone = timm.create_model(
            backbone_name, pretrained=False, num_classes=0,
            global_pool="avg", drop_rate=drop_rate)
        # probe feat_dim (mnv3/effb0 跟 num_features 不一致)
        with torch.no_grad():
            feat_dim = self.backbone(torch.zeros(1, 3, 64, 64)).shape[-1]
        self.feat_dim = feat_dim
        self.gender_head = nn.Linear(feat_dim, 1)
        self.age_head = nn.Linear(feat_dim, num_age)
    def forward(self, x):
        f = self.backbone(x)
        return self.gender_head(f).squeeze(-1), self.age_head(f)

# Load
ckpt = torch.load("age_gender_v20260507E_convnext_tiny.pt", weights_only=False)
# ckpt['args']['backbone'] = "convnext_tiny.fb_in22k_ft_in1k"
# ckpt['args']['img_h']=384, img_w=192
model = MultiHead(ckpt["args"]["backbone"]).eval()
model.load_state_dict(ckpt["model_state"])

# Inference (input: person crop)
mean = [0.485, 0.456, 0.406]; std = [0.229, 0.224, 0.225]
tf = T.Compose([T.Resize((384, 192)), T.ToTensor(), T.Normalize(mean, std)])

img = Image.open("person_crop.jpg").convert("RGB")
x = tf(img).unsqueeze(0)
with torch.no_grad():
    g_logit, a_logit = model(x)

gender_prob = torch.sigmoid(g_logit).item()      # > 0.5 = female
gender = "female" if gender_prob > 0.5 else "male"
age_idx = a_logit.argmax(dim=-1).item()
age_group = ["child", "young", "adult", "elder"][age_idx]

print(f"gender: {gender} ({gender_prob:.2f})")
print(f"age: {age_group}")
# 注意：young class 在 PA-100K/MSP60K 都沒有 supervised data，
# 所以 prediction 不會出 young；只會出 child/adult/elder