Update app.py

app.py

@@ -9,6 +9,15 @@ from detectron2.detectron2.config import get_cfg
 from detectron2 import model_zoo
 import torch_utils
 import dnnlib
+import torch
+import numpy as np
+import cv2
+import json
+import os
+import gradio as gr
+from detectron2.engine import DefaultPredictor
+from detectron2.config import get_cfg
+from detectron2 import model_zoo
 
 # Output directory
 output_dir = "key/"
@@ -30,7 +39,7 @@ def process_image(image, user_height_cm):
     keypoints = instances.pred_keypoints.cpu().numpy().tolist() if instances.has("pred_keypoints") else None
 
     if not keypoints:
-        return "No keypoints detected."
+        return "No keypoints detected.", None
 
     with open(output_file, "w") as f:
         json.dump({"keypoints": keypoints}, f, indent=4)
@@ -44,6 +53,15 @@ def process_image(image, user_height_cm):
     L_KNEE, R_KNEE = 13, 14
     L_ANKLE, R_ANKLE = 15, 16
 
+    skeleton = [(5, 6), (5, 11), (6, 12), (11, 12)]
+
+    for x, y in keypoints:
+        cv2.circle(image, (int(x), int(y)), 5, (0, 255, 0), -1)
+    for pt1, pt2 in skeleton:
+        x1, y1 = map(int, keypoints[pt1])
+        x2, y2 = map(int, keypoints[pt2])
+        cv2.line(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
+
     def get_distance(p1, p2):
         return np.linalg.norm(np.array(p1) - np.array(p2))
 
@@ -68,34 +86,30 @@ def process_image(image, user_height_cm):
     hip_circumference = waist_circumference / 0.75
 
     measurements = {
-        "Shoulder Width (cm)": round(shoulder_width_cm, 2),
-        "Waist Circumference (cm)": round(waist_circumference, 2),
-        "Hip Circumference (cm)": round(hip_circumference, 2),
-        "Torso Length (Neck to Pelvis, cm)": round(torso_length_cm, 2),
-        "Arm Length (Shoulder to Wrist, cm)": round(arm_length_cm, 2),
-        "Neck to Knee Length (cm)": round(neck_to_knee_cm, 2)
+        " Shoulder Width (cm)": round(shoulder_width_cm, 2),
+        " Waist Circumference (cm)": round(waist_circumference, 2),
+        " Hip Circumference (cm)": round(hip_circumference, 2),
+        " Torso Length (Neck to Pelvis, cm)": round(torso_length_cm, 2),
+        " Arm Length (Shoulder to Wrist, cm)": round(arm_length_cm, 2),
+        " Neck to Knee Length (cm)": round(neck_to_knee_cm, 2)
     }
 
-    return measurements
+    return measurements, cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
 # Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧍 AI-Powered Body Measurement Estimator")
-    gr.Markdown("Upload a clear **full-body image** and enter your **height (in cm)** to estimate key body dimensions using computer vision.")
+    gr.Markdown("## 🧍 Keypoint-Based Body Measurement Tool")
+    gr.Markdown("Upload a **full-body image** and enter your **height (in cm)** to estimate body measurements using AI-powered keypoint detection.")
 
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(type="pil", label="📸 Upload Image")
             height_input = gr.Number(label="📏 Your Height (cm)", value=170)
-            submit_btn = gr.Button("🔍 Estimate Measurements")
-
+            submit_btn = gr.Button("🔍 Generate Measurements")
         with gr.Column():
             measurement_output = gr.JSON(label="📐 Estimated Measurements")
+            image_output = gr.Image(type="pil", label="📌 Keypoint Overlay")
 
-    submit_btn.click(
-        fn=process_image,
-        inputs=[image_input, height_input],
-        outputs=[measurement_output]
-    )
+    submit_btn.click(fn=process_image, inputs=[image_input, height_input], outputs=[measurement_output, image_output])
 
 demo.launch()