Update app.py

app.py

@@ -2,6 +2,7 @@
 import io
 import os
 from typing import List
+
 import gradio as gr
 import torch
 import torch.nn as nn
@@ -10,59 +11,208 @@ import torchvision.transforms as T
 from PIL import Image, ImageDraw, ImageFont
 import numpy as np
 
-
+# Device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-
+# --- Model definition (must match your saved state) ---
 class AgeGenderClassifier(nn.Module):
-def __init__(self):
-super(AgeGenderClassifier, self).__init__()
-
-self.intermediate = nn.Sequential(
-nn.Linear(2048, 512),
-nn.ReLU(),
-nn.Dropout(0.4),
-nn.Linear(512, 128),
-nn.ReLU(),
-nn.Dropout(0.4),
-nn.Linear(128, 64),
-nn.ReLU(),
-)
-self.age_classifier = nn.Sequential(
-nn.Linear(64, 1),
-nn.Sigmoid()
-)
-self.gender_classifier = nn.Sequential(
-nn.Linear(64, 1),
-nn.Sigmoid()
-)
-
-
-def forward(self, x):
-x = self.intermediate(x)
-age = self.age_classifier(x)
-gender = self.gender_classifier(x)
-return age, gender
-
-
+    def __init__(self):
+        super(AgeGenderClassifier, self).__init__()
+        # classifier expected input dim 2048 (as in your training run)
+        self.intermediate = nn.Sequential(
+            nn.Linear(2048, 512),
+            nn.ReLU(),
+            nn.Dropout(0.4),
+            nn.Linear(512, 128),
+            nn.ReLU(),
+            nn.Dropout(0.4),
+            nn.Linear(128, 64),
+            nn.ReLU(),
+        )
+        self.age_classifier = nn.Sequential(
+            nn.Linear(64, 1),
+            nn.Sigmoid()
+        )
+        self.gender_classifier = nn.Sequential(
+            nn.Linear(64, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        x = self.intermediate(x)
+        age = self.age_classifier(x)
+        gender = self.gender_classifier(x)
+        return age, gender
 
 
 def build_model(weights_path: str):
-"""Rebuild VGG16 backbone + custom avgpool/classifier then load weights."""
-backbone = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1)
-
-for p in backbone.parameters():
-p.requires_grad = False
-
-
-
-for p in backbone.features[24:].parameters():
-p.requires_grad = True
-
-
-backbone.avgpool = nn.Sequential(
-nn.Conv2d(512, 512, kernel_size=3),
-nn.MaxPool2d(2),
-nn.ReLU(),
-nn.Flatten()
-demo.launch()
+    """Rebuild VGG16 backbone + custom avgpool/classifier then load weights."""
+    backbone = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1)
+    # freeze all then fine-tune later if needed (same as training script)
+    for p in backbone.parameters():
+        p.requires_grad = False
+
+    # allow last block to be trainable if desired (kept same as your training code)
+    for p in backbone.features[24:].parameters():
+        p.requires_grad = True
+
+    # replace avgpool with the same block used during training (conv->maxpool->relu->flatten)
+    backbone.avgpool = nn.Sequential(
+        nn.Conv2d(512, 512, kernel_size=3),
+        nn.MaxPool2d(2),
+        nn.ReLU(),
+        nn.Flatten()
+    )
+
+    # attach classifier
+    model = backbone
+    model.classifier = AgeGenderClassifier()
+
+    # load weights
+    if not os.path.exists(weights_path):
+        raise FileNotFoundError(f"Model weights not found at {weights_path}")
+
+    state = torch.load(weights_path, map_location=device)
+    # If saved state was model.state_dict(), load directly
+    try:
+        model.load_state_dict(state)
+    except Exception:
+        # if state is a dict with other keys, try common wrappers
+        if "model_state_dict" in state:
+            model.load_state_dict(state["model_state_dict"])
+        else:
+            raise
+
+    model.to(device)
+    model.eval()
+    return model
+
+
+# --- Preprocessing ---
+transform = T.Compose([
+    T.Resize((224, 224)),
+    T.ToTensor(),
+    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+INV_AGE_SCALE = 80  # training used age/80 normalization
+
+
+def draw_caption_on_image(pil_img: Image.Image, caption: str):
+    """Draw caption at the top of the image with a semi-transparent background."""
+    img = pil_img.convert("RGBA")
+    overlay = Image.new("RGBA", img.size, (255, 255, 255, 0))
+    draw = ImageDraw.Draw(overlay)
+
+    # choose a font size relative to image
+    fontsize = max(14, img.width // 20)
+    try:
+        font = ImageFont.truetype("DejaVuSans-Bold.ttf", fontsize)
+    except Exception:
+        font = ImageFont.load_default()
+
+    text_w, text_h = draw.textsize(caption, font=font)
+    padding = 8
+    rect_h = text_h + padding * 2
+
+    # draw translucent rectangle
+    draw.rectangle([(0, 0), (img.width, rect_h)], fill=(0, 0, 0, 160))
+    # write text
+    draw.text((padding, padding), caption, font=font, fill=(255, 255, 255, 255))
+
+    out = Image.alpha_composite(img, overlay).convert("RGB")
+    return out
+
+
+# --- Prediction function for multiple images ---
+
+def predict_images(images: List[Image.Image], model) -> List[Image.Image]:
+    """Takes a list of PIL images and returns list of PIL images annotated with predictions."""
+    if images is None or len(images) == 0:
+        return []
+
+    # preprocess all images into a batch
+    tensors = []
+    for im in images:
+        if im.mode != "RGB":
+            im = im.convert("RGB")
+        t = transform(im)
+        tensors.append(t)
+
+    batch = torch.stack(tensors).to(device)
+
+    with torch.no_grad():
+        pred_age, pred_gender = model(batch)
+        # ensure shapes (N,1)
+        pred_age = pred_age.squeeze(-1).cpu().numpy()
+        pred_gender = pred_gender.squeeze(-1).cpu().numpy()
+
+    outputs = []
+    for img, pa, pg in zip(images, pred_age, pred_gender):
+        age_val = int(np.clip(pa, 0.0, 1.0) * INV_AGE_SCALE)
+        gender_label = "Female" if pg > 0.5 else "Male"
+        gender_emoji = "👩" if pg > 0.5 else "👨"
+        conf = float(pg if pg > 0.5 else 1 - pg)
+
+        caption = f"{gender_emoji} {gender_label} ({conf:.2f})  •  🎂 Age ≈ {age_val}"
+        out_img = draw_caption_on_image(img, caption)
+        outputs.append(out_img)
+
+    return outputs
+
+
+# --- Load model once on startup ---
+MODEL_WEIGHTS = os.environ.get("MODEL_PATH", "age_gender_model.pth")
+model = build_model(MODEL_WEIGHTS)
+
+# --- Gradio UI ---
+with gr.Blocks(title="FairFace Age & Gender — Multi-image Demo") as demo:
+    gr.Markdown("""
+    # 🧠 FairFace Multi-task Age & Gender Predictor
+    Upload **one or more** images (JPG/PNG). The app will predict **gender** and **age** for each image and display results right on the picture.
+
+    **How to use**
+    1. Click **Browse** or drag & drop multiple images. ✅
+    2. Click **Run**. The model processes images and shows results below. ⚡
+    3. Use the download button on the output images if you want to save them.
+
+    *Note:* Age is estimated (approx.). This model was trained on the FairFace dataset.
+    """)
+
+    with gr.Row():
+        img_input = gr.File(file_count="multiple", label="Upload images")
+        run_btn = gr.Button("Run  ▶️")
+
+    gallery = gr.Gallery(label="Predictions", show_label=True, elem_id="gallery").style(grid=[3], height="auto")
+
+    def run_and_predict(files):
+        # files is list of uploaded file dicts or file paths depending on environment
+        if not files:
+            return []
+
+        pil_imgs = []
+        # if File component returns list of dicts in HF spaces, handle both
+        for f in files:
+            # f might be a path string or dict-like
+            if isinstance(f, dict) and "name" in f and "data" in f:
+                # web upload format
+                im = Image.open(io.BytesIO(f["data"]))
+            else:
+                path = f if isinstance(f, str) else f.name
+                im = Image.open(path)
+            pil_imgs.append(im.convert("RGB"))
+
+        return predict_images(pil_imgs, model)
+
+    run_btn.click(fn=run_and_predict, inputs=[img_input], outputs=[gallery])
+
+    gr.Markdown("""
+    ---
+    **Tips & Notes**
+    - The model outputs age normalized to 0–80 years (approx).
+    - If results look odd, try a clearer, frontal face image.
+    - This demo is for research / demo purposes only — be mindful of privacy. 🙏
+    """)
+
+if __name__ == "__main__":
+    demo.launch()