Update app.py

app.py

@@ -8,17 +8,16 @@ import torch
 import torch.nn as nn
 import torchvision.models as models
 import torchvision.transforms as T
-from PIL import ImageFont, ImageDraw, Image
+from PIL import Image
 import numpy as np
 
 # Device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# --- Model definition (must match your saved state) ---
+# --- Model definition ---
 class AgeGenderClassifier(nn.Module):
     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(),
@@ -48,15 +47,15 @@ class AgeGenderClassifier(nn.Module):
 def build_model(weights_path: str):
     """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)
+
+    # freeze all layers
     for p in backbone.parameters():
         p.requires_grad = False
-
-    # allow last block to be trainable if desired (kept same as your training code)
+    # optionally allow last block to be trainable
     for p in backbone.features[24:].parameters():
         p.requires_grad = True
 
-    # replace avgpool with the same block used during training (conv->maxpool->relu->flatten)
+    # replace avgpool
     backbone.avgpool = nn.Sequential(
         nn.Conv2d(512, 512, kernel_size=3),
         nn.MaxPool2d(2),
@@ -73,11 +72,9 @@ def build_model(weights_path: str):
         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:
@@ -98,81 +95,50 @@ transform = T.Compose([
 INV_AGE_SCALE = 80  # training used age/80 normalization
 
 
-def draw_caption_on_image(image, caption):
-    if image.mode != "RGBA":
-        image = image.convert("RGBA")
-    draw = ImageDraw.Draw(image)
-    font = ImageFont.load_default()
-    
-    bbox = draw.textbbox((0,0), caption, font=font)
-    text_w = bbox[2] - bbox[0]
-    text_h = bbox[3] - bbox[1]
-
-    # semi-transparent rectangle
-    overlay = Image.new("RGBA", image.size)
-    overlay_draw = ImageDraw.Draw(overlay)
-    overlay_draw.rectangle([0,0,text_w+20,text_h+20], fill=(0,0,0,127))
-    image = Image.alpha_composite(image, overlay)
-
-    draw = ImageDraw.Draw(image)
-    draw.text((10,10), caption, font=font, fill="white")
-    return image.convert("RGB")
-
-
-
-# --- Prediction function for multiple images ---
+# --- Prediction function ---
+def predict_images_with_text(images: List[Image.Image], model):
+    """Return original images and captions for each."""
+    if not images:
+        return [], []
 
-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)
+        tensors.append(transform(im))
 
     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 = []
+    output_images = []
+    captions = []
+
     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)
+        output_images.append(np.array(img))
+        captions.append(f"{gender_emoji} {gender_label} ({conf:.2f})  •  🎂 Age ≈ {age_val}")
 
-    return outputs
+    return output_images, captions
 
 
-# --- Load model once on startup ---
+# --- Load model ---
 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.
+    Upload **one or more** images (JPG/PNG). The app will predict **gender** and **age** for each image and display results below the image.
     """)
 
     with gr.Row():
@@ -180,41 +146,34 @@ with gr.Blocks(title="FairFace Age & Gender — Multi-image Demo") as demo:
         run_btn = gr.Button("Run  ▶️")
 
     gallery = gr.Gallery(
-        label="Predictions",
-        show_label=True,
-        elem_id="gallery",
-        columns=3,   # 3 images per row
+        label="Uploaded Images",
+        columns=3,
         height="auto"
     )
 
+    captions = gr.HTML(label="Predictions")
+
     def run_and_predict(files):
-        # files is list of uploaded file dicts or file paths depending on environment
         if not files:
-            return []
-
+            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"))
+            path = f if isinstance(f, str) else f.name
+            pil_imgs.append(Image.open(path).convert("RGB"))
 
-        return predict_images(pil_imgs, model)
+        imgs, texts = predict_images_with_text(pil_imgs, model)
+        captions_html = "<br>".join([f"<h2>{t}</h2>" for t in texts])
+        return imgs, captions_html
 
-    run_btn.click(fn=run_and_predict, inputs=[img_input], outputs=[gallery])
+    run_btn.click(fn=run_and_predict, inputs=[img_input], outputs=[gallery, captions])
 
     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. 🙏
+    - Age is normalized to 0–80 years (approx.).
+    - For best results, upload clear frontal face images.
+    - This is a demo — respect privacy when using photos. 🙏
     """)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()