Update app.py

app.py

@@ -508,6 +508,7 @@
 
 # if __name__ == "__main__":
 #     demo.launch()
+
 import torch
 import numpy as np
 import cv2
@@ -518,7 +519,7 @@ from detectron2.engine import DefaultPredictor
 from detectron2.config import get_cfg
 from detectron2 import model_zoo
 
-# Create output directory if it doesn't exist
+# Output directory
 output_dir = "key/"
 os.makedirs(output_dir, exist_ok=True)
 output_file = os.path.join(output_dir, "keypoints.json")
@@ -528,106 +529,86 @@ cfg = get_cfg()
 cfg.merge_from_file(model_zoo.get_config_file("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml"))
 cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml")
 cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load the predictor
 predictor = DefaultPredictor(cfg)
 
 def process_image(image, user_height_cm):
-    # Convert Gradio image input to OpenCV format
     image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-
-    # Run keypoint detection
     outputs = predictor(image)
 
-    # Extract keypoints
     instances = outputs["instances"]
     keypoints = instances.pred_keypoints.cpu().numpy().tolist() if instances.has("pred_keypoints") else None
 
     if not keypoints:
         return "No keypoints detected.", None
 
-    # Save keypoints to JSON
     with open(output_file, "w") as f:
         json.dump({"keypoints": keypoints}, f, indent=4)
 
-    keypoints = np.array(keypoints[0])[:, :2]  # Extract (x, y) coordinates
+    keypoints = np.array(keypoints[0])[:, :2]
 
-    # COCO format indices
     NOSE, L_SHOULDER, R_SHOULDER = 0, 5, 6
     L_ELBOW, R_ELBOW = 7, 8
     L_WRIST, R_WRIST = 9, 10
     L_HIP, R_HIP = 11, 12
+    L_KNEE, R_KNEE = 13, 14
     L_ANKLE, R_ANKLE = 15, 16
 
-    # Define Keypoint Pairs for Drawing Lines (COCO Format)
     skeleton = [(5, 6), (5, 11), (6, 12), (11, 12)]
 
-    # Draw Keypoints
     for x, y in keypoints:
         cv2.circle(image, (int(x), int(y)), 5, (0, 255, 0), -1)
-
-    # Draw Skeleton
     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)
 
-    # Function to calculate Euclidean distance
     def get_distance(p1, p2):
         return np.linalg.norm(np.array(p1) - np.array(p2))
 
-    # Calculate full height (consider head length)
     ankle_mid = ((keypoints[L_ANKLE] + keypoints[R_ANKLE]) / 2).tolist()
     pixel_height = get_distance(keypoints[NOSE], ankle_mid)
-
-    # Estimated full body height (add approx head length)
-    estimated_full_pixel_height = pixel_height / 0.87  # Since 87% = nose to ankle
+    estimated_full_pixel_height = pixel_height / 0.87
     pixels_per_cm = estimated_full_pixel_height / user_height_cm
 
-    # Waist and shoulder measurements
-    shoulder_width_px = get_distance(keypoints[L_SHOULDER], keypoints[R_SHOULDER])
-    waist_width_px = get_distance(keypoints[L_HIP], keypoints[R_HIP])
-
-    # Convert to cm
-    shoulder_width_cm = shoulder_width_px / pixels_per_cm
-    waist_width_cm = waist_width_px / pixels_per_cm
+    shoulder_width_cm = get_distance(keypoints[L_SHOULDER], keypoints[R_SHOULDER]) / pixels_per_cm
+    waist_width_cm = get_distance(keypoints[L_HIP], keypoints[R_HIP]) / pixels_per_cm
 
-    # Torso Length (Neck to Pelvis)
     pelvis = ((keypoints[L_HIP] + keypoints[R_HIP]) / 2).tolist()
     neck = ((keypoints[L_SHOULDER] + keypoints[R_SHOULDER]) / 2).tolist()
-    torso_length_px = get_distance(neck, pelvis)
-    torso_length_cm = torso_length_px / pixels_per_cm
+    torso_length_cm = get_distance(neck, pelvis) / pixels_per_cm
 
-    # Arm Length (Shoulder to Wrist)
-    arm_length_px = get_distance(keypoints[L_SHOULDER], keypoints[L_WRIST])
-    arm_length_cm = arm_length_px / pixels_per_cm
+    arm_length_cm = get_distance(keypoints[L_SHOULDER], keypoints[L_WRIST]) / pixels_per_cm
 
-    # Calculate waist and hip circumference (Ellipse approximation)
-    # Waist circumference ≈ π × (waist_width / 2) × 2
-    waist_circumference = np.pi * waist_width_cm
-    hip_circumference = waist_circumference / 0.75  # Assuming hip is slightly bigger than waist
+    knee_mid = ((keypoints[L_KNEE] + keypoints[R_KNEE]) / 2).tolist()
+    neck_to_knee_cm = get_distance(neck, knee_mid) / pixels_per_cm
 
-    # Improved body measurement calculation
-    def calculate_body_measurements(waist_circumference, hip_circumference, shoulder_width_cm, torso_length_cm, arm_length_cm):
-        return {
-            "Waist Circumference (cm)": round(waist_circumference, 2),
-            "Hip Circumference (cm)": round(hip_circumference, 2),
-            "Shoulder Width (cm)": round(shoulder_width_cm, 2),
-            "Torso Length (Neck to Pelvis, cm)": round(torso_length_cm, 2),
-            "Full Arm Length (Shoulder to Wrist, cm)": round(arm_length_cm, 2),
-        }
+    waist_circumference = np.pi * waist_width_cm
+    hip_circumference = waist_circumference / 0.75
 
-    measurements = calculate_body_measurements(waist_circumference, hip_circumference, shoulder_width_cm, torso_length_cm, arm_length_cm)
+    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)
+    }
 
     return measurements, cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
 # Gradio Interface
-demo = gr.Interface(
-    fn=process_image,
-    inputs=[gr.Image(type="pil"), gr.Number(label="User Height (cm)")],
-    outputs=[gr.JSON(label="Measurements"), gr.Image(type="pil", label="Keypoint Overlay")],
-    title="Keypoint Measurement Extractor",
-    description="Upload an image, enter your height, and get body measurements based on keypoints.",
-)
+with gr.Blocks() as demo:
+    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")
+            submit_btn = gr.Button("🔍 Generate Measurements")
+        with gr.Column():
+            height_input = gr.Number(label="📏 Your Height (cm)", value=170)
+            measurement_output = gr.JSON(label="📐 Estimated Measurements")
+
+    submit_btn.click(fn=process_image, inputs=[image_input, height_input], outputs=[measurement_output, image_output])
 
 demo.launch()