implement UI and basic functions

app.py

@@ -5,11 +5,12 @@ 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]
+human_list = os.listdir(os.path.join(example_path, "human"))
+human_list_path = [os.path.join(example_path, "human", human) for human in human_list]
 
 # Initialize MediaPipe Pose
 mp_pose = mp.solutions.pose
@@ -18,76 +19,65 @@ mp_drawing = mp.solutions.drawing_utils
 mp_pose_landmark = mp_pose.PoseLandmark
 
 
-def align_clothing(body_img, clothing_img):
-    image_rgb = cv2.cvtColor(body_img, cv2.COLOR_BGR2RGB)
+def detect_pose(image):
+    # Convert to RGB
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+    # Run pose detection
     result = pose.process(image_rgb)
-    output = body_img.copy()
+
     keypoints = {}
 
     if result.pose_landmarks:
-        height, width, _ = output.shape
+        # Draw landmarks on image
+        mp_drawing.draw_landmarks(image, result.pose_landmarks, mp_pose.POSE_CONNECTIONS)
+
+        # Get image dimensions
+        height, width, _ = image.shape
 
-        # Extract body keypoints
-        points = {
+        # Extract specific landmarks
+        landmark_indices = {
             'left_shoulder': mp_pose_landmark.LEFT_SHOULDER,
             'right_shoulder': mp_pose_landmark.RIGHT_SHOULDER,
-            'left_hip': mp_pose_landmark.LEFT_HIP
+            'left_hip': mp_pose_landmark.LEFT_HIP,
+            'right_hip': mp_pose_landmark.RIGHT_HIP
         }
 
-        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()
+        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
+
+
+def process_image(human_img):
+    # Convert PIL image to NumPy array
+    human_img = np.array(human_img)
+
+    processed_image = detect_pose(human_img)
+    return processed_image
+
+
+image_blocks = gr.Blocks().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)
+            human_img = gr.Image(type="pil", label='Human', interactive=True)
             example = gr.Examples(
-                inputs=imgs,
+                inputs=human_img,
                 examples_per_page=10,
                 examples=human_list_path
             )
 
         with gr.Column():
-            garm_img = gr.Image(label="Garment", type="pil",interactive=True)
+            garm_img = gr.Image(label="Garment", type="pil", interactive=True)
             example = gr.Examples(
                 inputs=garm_img,
                 examples_per_page=8,
@@ -96,5 +86,9 @@ with image_blocks as demo:
             image_out = gr.Image(label="Processed image", type="pil")
 
     with gr.Row():
-        try_button = gr.Button(value="Try-on")
+        try_button = gr.Button(value="Try-on", variant='primary')
+
+    # Linking the button to the processing function
+    try_button.click(fn=process_image, inputs=human_img, outputs=image_out)
+
 image_blocks.launch()