Update app.py

app.py

@@ -4,10 +4,10 @@ import torch.nn as nn
 import torchvision.transforms as transforms
 import numpy as np
 from PIL import Image
-import mediapipe as mp
+import cv2
 import os
 
-# ── Model Definition (must match training exactly) ────────────────────────────
+# ── Model Definition ──────────────────────────────────────────────────────────
 class ConvBlock(nn.Module):
     def __init__(self, in_ch, out_ch):
         super().__init__()
@@ -20,7 +20,8 @@ class ConvBlock(nn.Module):
     def forward(self, x): return self.block(x)
 
 class TryOnUNet(nn.Module):
-    def __init__(self, in_ch=25, base_ch=32):
+    def __init__(self, in_ch=22, base_ch=32):
+        # in_ch = agnostic(3) + cloth(3) + cloth_mask(1) + pose_heatmap(15)
         super().__init__()
         self.enc1 = ConvBlock(in_ch, base_ch)
         self.enc2 = ConvBlock(base_ch, base_ch*2)
@@ -50,41 +51,77 @@ H, W   = 256, 192
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 # ── Load Model ────────────────────────────────────────────────────────────────
-def load_model(weights_path='tryon_model.pth'):
-    model = TryOnUNet(in_ch=25, base_ch=32).to(DEVICE)
-    if os.path.exists(weights_path):
-        ckpt = torch.load(weights_path, map_location=DEVICE)
-        # Handle both raw state_dict and checkpoint dict
-        sd = ckpt.get('model_state_dict', ckpt)
-        model.load_state_dict(sd)
-        print(f"✅ Loaded weights from {weights_path}")
+model = TryOnUNet(in_ch=22, base_ch=32).to(DEVICE)
+if os.path.exists('tryon_model.pth'):
+    ckpt = torch.load('tryon_model.pth', map_location=DEVICE)
+    sd   = ckpt.get('model_state_dict', ckpt)
+    model.load_state_dict(sd, strict=False)
+    print("✅ Weights loaded")
+else:
+    print("⚠️  No weights — using random weights (demo)")
+model.eval()
+
+# ── Pose: OpenCV-based body keypoint approximation ────────────────────────────
+# Uses HOG person detector + simple torso geometry
+# No external pose library needed — works on any Python 3.13 environment
+
+def estimate_pose_heatmap(img_pil, h=H, w=W):
+    """
+    Lightweight pose heatmap using OpenCV HOG detector.
+    Approximates 15 keypoint locations from detected person bounding box.
+    Good enough for clothing region alignment in try-on.
+    """
+    img_np = np.array(img_pil.convert('RGB'))
+    img_resized = cv2.resize(img_np, (w, h))
+    gray = cv2.cvtColor(img_resized, cv2.COLOR_RGB2GRAY)
+
+    heatmap = np.zeros((15, h, w), dtype=np.float32)
+
+    # Try HOG person detection
+    hog = cv2.HOGDescriptor()
+    hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())
+    img_bgr = cv2.cvtColor(img_resized, cv2.COLOR_RGB2BGR)
+    boxes, weights = hog.detectMultiScale(img_bgr, winStride=(8,8), padding=(4,4), scale=1.05)
+
+    if len(boxes) > 0:
+        # Use largest detected person box
+        idx = np.argmax([b[2]*b[3] for b in boxes])
+        x, y, bw, bh = boxes[idx]
+        cx = x + bw // 2   # center x
     else:
-        print("⚠️  No weights file found — using random weights (demo only)")
-    model.eval()
-    return model
-
-model = load_model()
-
-# ── Pose Estimator ────────────────────────────────────────────────────────────
-mp_pose = mp.solutions.pose
-pose_estimator = mp_pose.Pose(
-    static_image_mode=True, model_complexity=1,
-    min_detection_confidence=0.5
-)
+        # Fallback: assume person is centered
+        x, y, bw, bh = w//4, 0, w//2, h
+        cx = w // 2
+
+    # Approximate keypoint positions from bounding box geometry
+    # Order: nose, neck, L-shoulder, R-shoulder, L-elbow, R-elbow,
+    #        L-wrist, R-wrist, L-hip, R-hip, L-knee, R-knee,
+    #        L-ankle, R-ankle, center-chest
+    keypoints = [
+        (cx,          y + int(bh*0.05)),   # 0  nose
+        (cx,          y + int(bh*0.12)),   # 1  neck
+        (cx - bw//4,  y + int(bh*0.18)),   # 2  left shoulder
+        (cx + bw//4,  y + int(bh*0.18)),   # 3  right shoulder
+        (cx - bw//3,  y + int(bh*0.35)),   # 4  left elbow
+        (cx + bw//3,  y + int(bh*0.35)),   # 5  right elbow
+        (cx - bw//3,  y + int(bh*0.52)),   # 6  left wrist
+        (cx + bw//3,  y + int(bh*0.52)),   # 7  right wrist
+        (cx - bw//5,  y + int(bh*0.55)),   # 8  left hip
+        (cx + bw//5,  y + int(bh*0.55)),   # 9  right hip
+        (cx - bw//5,  y + int(bh*0.72)),   # 10 left knee
+        (cx + bw//5,  y + int(bh*0.72)),   # 11 right knee
+        (cx - bw//5,  y + int(bh*0.90)),   # 12 left ankle
+        (cx + bw//5,  y + int(bh*0.90)),   # 13 right ankle
+        (cx,          y + int(bh*0.30)),   # 14 chest center
+    ]
+
+    sigma = 10
+    yy, xx = np.ogrid[:h, :w]
+    for i, (kx, ky) in enumerate(keypoints):
+        kx = int(np.clip(kx, 0, w-1))
+        ky = int(np.clip(ky, 0, h-1))
+        heatmap[i] = np.exp(-((xx-kx)**2 + (yy-ky)**2) / (2*sigma**2))
 
-def extract_pose_heatmap(img_pil):
-    img_np = np.array(img_pil.convert('RGB'))
-    results = pose_estimator.process(img_np)
-    heatmap = np.zeros((18, H, W), dtype=np.float32)
-    if results.pose_landmarks:
-        lms = results.pose_landmarks.landmark
-        for i in range(min(18, len(lms))):
-            x = int(lms[i].x * W)
-            y = int(lms[i].y * H)
-            c = lms[i].visibility
-            if c > 0.3 and 0 <= x < W and 0 <= y < H:
-                yy, xx = np.ogrid[:H, :W]
-                heatmap[i] = np.exp(-((xx-x)**2+(yy-y)**2)/(2*8**2)) * c
     return torch.from_numpy(heatmap)
 
 # ── Transforms ────────────────────────────────────────────────────────────────
@@ -99,32 +136,25 @@ def denorm(tensor):
     return transforms.ToPILImage()(t.clamp(0, 1))
 
 def make_agnostic(person_t):
+    """Mask out torso clothing region."""
     ag = person_t.clone()
     ag[:, H//4:3*H//4, W//6:5*W//6] = -0.5
     return ag
 
-# ── Main Inference Function ───────────────────────────────────────────────────
+# ── Inference ─────────────────────────────────────────────────────────────────
 def virtual_tryon(person_img, cloth_img):
-    """
-    Args:
-        person_img : PIL Image — full body photo of person
-        cloth_img  : PIL Image — flat-lay garment photo
-    Returns:
-        result_img : PIL Image — person wearing the garment
-    """
     if person_img is None or cloth_img is None:
-        return None, "⚠️ Please upload both a person photo and a garment photo."
-
+        gr.Warning("Please upload both a person photo and a garment photo.")
+        return None
     try:
-        # Preprocess
-        person_t = img_transform(person_img.convert('RGB'))
-        cloth_t  = img_transform(cloth_img.convert('RGB'))
-        agnostic = make_agnostic(person_t)
-        pose_map = extract_pose_heatmap(person_img)
+        person_t   = img_transform(person_img.convert('RGB'))
+        cloth_t    = img_transform(cloth_img.convert('RGB'))
+        agnostic   = make_agnostic(person_t)
+        pose_map   = estimate_pose_heatmap(person_img)   # [15, H, W]
         cloth_mask = torch.zeros(1, H, W)
         cloth_mask[:, H//4:3*H//4, W//6:5*W//6] = 1.0
 
-        # Stack inputs → [1, 25, H, W]
+        # Stack: 3 + 3 + 1 + 15 = 22 channels
         inp = torch.cat([
             agnostic.unsqueeze(0),
             cloth_t.unsqueeze(0),
@@ -132,86 +162,41 @@ def virtual_tryon(person_img, cloth_img):
             pose_map.unsqueeze(0)
         ], dim=1).to(DEVICE)
 
-        # Inference
         with torch.no_grad():
             out = model(inp).squeeze(0)
 
-        result = denorm(out)
-        return result, "✅ Try-on complete!"
+        return denorm(out)
 
     except Exception as e:
-        return None, f"❌ Error: {str(e)}"
+        raise gr.Error(f"Inference failed: {str(e)}")
 
 # ── Gradio UI ─────────────────────────────────────────────────────────────────
-with gr.Blocks(
-    title="Virtual Try-On",
-    theme=gr.themes.Soft(),
-    css="""
-        .header { text-align: center; padding: 20px 0 10px; }
-        .header h1 { font-size: 2rem; font-weight: 700; color: #1a1a2e; }
-        .header p  { color: #555; font-size: 1rem; }
-        .result-box { border: 2px dashed #ccc; border-radius: 12px; }
-        footer { display: none !important; }
-    """
-) as demo:
-
-    gr.HTML("""
-        <div class='header'>
-            <h1>🧥 Virtual Try-On System</h1>
-            <p>Upload a person photo and a garment — see it on them instantly</p>
-            <p style='font-size:0.85rem; color:#888;'>
-                FYP Project | BZU CASPAM | Post ADP Mathematics
-            </p>
-        </div>
+with gr.Blocks(title="Virtual Try-On", theme=gr.themes.Soft()) as demo:
+
+    gr.Markdown("""
+    # 🧥 Virtual Try-On System
+    Upload a **person photo** and a **garment image** to see the outfit on them.
+    > FYP Project · BZU CASPAM · Post ADP Mathematics
     """)
 
     with gr.Row():
-        with gr.Column(scale=1):
-            person_input = gr.Image(
-                label="👤 Person Photo",
-                type="pil",
-                height=320
-            )
-            gr.Markdown("> Full-body frontal photo works best")
-
-        with gr.Column(scale=1):
-            cloth_input = gr.Image(
-                label="👕 Garment Photo",
-                type="pil",
-                height=320
-            )
-            gr.Markdown("> Flat-lay product photo on white background")
-
-        with gr.Column(scale=1):
-            result_output = gr.Image(
-                label="✨ Try-On Result",
-                type="pil",
-                height=320,
-                elem_classes=["result-box"]
-            )
-            status_text = gr.Markdown("")
+        person_input = gr.Image(label="👤 Person Photo",  type="pil", height=320)
+        cloth_input  = gr.Image(label="👕 Garment Photo", type="pil", height=320)
+        result_img   = gr.Image(label="✨ Try-On Result", type="pil", height=320)
 
     run_btn = gr.Button("🚀 Try It On!", variant="primary", size="lg")
 
+    gr.Markdown("""
+    **Tips for best results:**
+    - Use a **frontal full-body** photo with plain background
+    - Garment should be **flat-lay on white background**
+    - Works best with **tops and shirts**
+    """)
+
     run_btn.click(
         fn=virtual_tryon,
         inputs=[person_input, cloth_input],
-        outputs=[result_output, status_text]
+        outputs=[result_img]
     )
 
-    gr.Markdown("""
-    ---
-    ### 📌 Tips for best results
-    - Use a **frontal, full-body** photo with clear background
-    - Garment should be a **flat-lay on white background**
-    - Works best with **tops and shirts** (not pants/shoes yet)
-
-    ### 🔧 How it works
-    1. **Pose Estimation** — MediaPipe detects 18 body keypoints
-    2. **Agnostic Masking** — Original clothing region is masked out
-    3. **UNet Synthesis** — Model warps garment onto body shape
-    """)
-
-if __name__ == "__main__":
-    demo.launch()
-    
+demo.launch()