edit UI, add logic to detect keypoints

app.py

@@ -2,6 +2,14 @@ import gradio as gr
 import cv2
 import numpy as np
 import mediapipe as mp
+import os
+
+example_path = os.path.join(os.path.dirname(__file__), 'example')
+garm_list = os.listdir(os.path.join(example_path, "cloth"))
+garm_list_path = [os.path.join(example_path, "cloth", garm) for garm in garm_list]
+
+human_list = os.listdir(os.path.join(example_path, "cloth"))
+human_list_path = [os.path.join(example_path, "cloth", garm) for garm in garm_list]
 
 # Initialize MediaPipe Pose
 mp_pose = mp.solutions.pose
@@ -9,52 +17,84 @@ pose = mp_pose.Pose(static_image_mode=True)
 mp_drawing = mp.solutions.drawing_utils
 mp_pose_landmark = mp_pose.PoseLandmark
 
-def detect_pose(image):
-    # Convert to RGB
-    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
-    # Run pose detection
+def align_clothing(body_img, clothing_img):
+    image_rgb = cv2.cvtColor(body_img, cv2.COLOR_BGR2RGB)
     result = pose.process(image_rgb)
-
+    output = body_img.copy()
     keypoints = {}
 
     if result.pose_landmarks:
-        # Draw landmarks on image
-        mp_drawing.draw_landmarks(image, result.pose_landmarks, mp_pose.POSE_CONNECTIONS)
-
-        # Get image dimensions
-        height, width, _ = image.shape
+        height, width, _ = output.shape
 
-        # Extract specific landmarks
-        landmark_indices = {
+        # Extract body keypoints
+        points = {
             'left_shoulder': mp_pose_landmark.LEFT_SHOULDER,
             'right_shoulder': mp_pose_landmark.RIGHT_SHOULDER,
-            'left_hip': mp_pose_landmark.LEFT_HIP,
-            'right_hip': mp_pose_landmark.RIGHT_HIP
+            'left_hip': mp_pose_landmark.LEFT_HIP
         }
 
-        for name, index in landmark_indices.items():
-            lm = result.pose_landmarks.landmark[index]
-            x, y = int(lm.x * width), int(lm.y * height)
-            keypoints[name] = (x, y)
-
-            # Draw a circle + label for debug
-            cv2.circle(image, (x, y), 5, (0, 255, 0), -1)
-            cv2.putText(image, name, (x + 5, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
-
-    return image, keypoints
-
-# Gradio interface
-iface = gr.Interface(
-    fn=detect_pose,
-    inputs=gr.Image(type="numpy", label="Upload Full-Body Image"),
-    outputs=[
-        gr.Image(type="numpy", label="Pose Visualization"),
-        gr.JSON(label="Extracted Keypoints")
-    ],
-    title="Virtual Try-On - Pose Detection",
-    description="Detects body keypoints using MediaPipe Pose and visualizes them. Shoulders and hips are labeled."
-)
-
-if __name__ == "__main__":
-    iface.launch()
+        for name, idx in points.items():
+            lm = result.pose_landmarks.landmark[idx]
+            keypoints[name] = (int(lm.x * width), int(lm.y * height))
+
+        # Draw for debug
+        for name, (x, y) in keypoints.items():
+            cv2.circle(output, (x, y), 5, (0, 255, 0), -1)
+            cv2.putText(output, name, (x + 5, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
+
+        # Affine Transform
+        if all(k in keypoints for k in ['left_shoulder', 'right_shoulder', 'left_hip']):
+            src_tri = np.array([
+                [0, 0],
+                [clothing_img.shape[1], 0],
+                [0, clothing_img.shape[0]]
+            ], dtype=np.float32)
+
+            dst_tri = np.array([
+                keypoints['left_shoulder'],
+                keypoints['right_shoulder'],
+                keypoints['left_hip']
+            ], dtype=np.float32)
+
+            # Compute warp matrix and apply it
+            warp_mat = cv2.getAffineTransform(src_tri, dst_tri)
+            warped_clothing = cv2.warpAffine(clothing_img, warp_mat, (width, height), flags=cv2.INTER_LINEAR,
+                                             borderMode=cv2.BORDER_TRANSPARENT)
+
+            # Blend clothing over body
+            if clothing_img.shape[2] == 4:  # has alpha
+                alpha = warped_clothing[:, :, 3] / 255.0
+                for c in range(3):
+                    output[:, :, c] = (1 - alpha) * output[:, :, c] + alpha * warped_clothing[:, :, c]
+            else:
+                output = cv2.addWeighted(output, 0.8, warped_clothing, 0.5, 0)
+
+    return output
+
+
+image_blocks = gr.Blocks(theme="Nymbo/Alyx_Theme").queue()
+with image_blocks as demo:
+    gr.HTML("<center><h1>Virtual Try-On</h1></center>")
+    gr.HTML("<center><p>Upload an image of a person and an image of a garment ✨</p></center>")
+    with gr.Row():
+        with gr.Column():
+            imgs = gr.Image(type="pil", label='Human', interactive=True)
+            example = gr.Examples(
+                inputs=imgs,
+                examples_per_page=10,
+                examples=human_list_path
+            )
+
+        with gr.Column():
+            garm_img = gr.Image(label="Garment", type="pil",interactive=True)
+            example = gr.Examples(
+                inputs=garm_img,
+                examples_per_page=8,
+                examples=garm_list_path)
+        with gr.Column():
+            image_out = gr.Image(label="Processed image", type="pil")
+
+    with gr.Row():
+        try_button = gr.Button(value="Try-on")
+image_blocks.launch()