Update app.py

app.py

@@ -118,158 +118,183 @@
 
 
 
-
+import gradio as gr
+import json
 import base64
-from io import BytesIO
-from PIL import Image
 import requests
+from io import BytesIO
 import torch
 import numpy as np
 import cv2
-import json
-import os
-import gradio as gr
+from PIL import Image
 from detectron2.engine import DefaultPredictor
 from detectron2.config import get_cfg
 from detectron2 import model_zoo
 
-# === Save to DB function ===
-def save_to_database(measurements, image, user_height_cm, user_id):
-    if not user_id:
-        return "❌ user_id missing from URL."
-    if measurements is None or image is None:
-        return "⚠️ No data to save."
-
-    buffered = BytesIO()
-    pil_image = Image.fromarray(image)
-    pil_image.save(buffered, format="JPEG")
-    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-
-    payload = {
-        "imageBase64": img_str,
-        "heightCm": user_height_cm,
-        "waistCircumferenceCm": measurements.get("Waist Circumference (cm)"),
-        "hipCircumferenceCm": measurements.get("Hip Circumference (cm)"),
-        "shoulderWidthCm": measurements.get("Shoulder Width (cm)"),
-        "torsoLengthCm": measurements.get("Torso Length (Neck to Pelvis, cm)"),
-        "fullArmLengthCm": measurements.get("Arm Length (Shoulder to Wrist, cm)"),
-        "neckToKneeLengthCm": measurements.get("Neck to Knee Length (cm)"),
-    }
-
-    try:
-        response = requests.post(
-            f"https://7da2-2409-4042-6e81-1806-de6-b8e5-836c-6b95.ngrok-free.app/upload/{user_id}",
-            json=payload
-        )
-        if response.status_code == 201:
-            return "✅ Measurements and image saved to database!"
-        else:
-            return f"❌ Failed: {response.status_code} - {response.text}"
-    except Exception as e:
-        return f"⚠️ Error during save: {str(e)}"
-
-# Setup Detectron2 Keypoint R-CNN model
-output_dir = "key/"
-os.makedirs(output_dir, exist_ok=True)
-output_file = os.path.join(output_dir, "keypoints.json")
-
+# === Set up Detectron2 model ===
 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"
 predictor = DefaultPredictor(cfg)
 
-def process_image(image, user_height_cm):
-    image_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-    outputs = predictor(image_np)
+# === Utility ===
+def get_distance(p1, p2):
+    return np.linalg.norm(np.array(p1) - np.array(p2))
 
+# === Keypoint and Measurement Logic ===
+def process_image(image, user_height_cm):
+    image_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    outputs = predictor(image_cv)
     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
-
-    with open(output_file, "w") as f:
-        json.dump({"keypoints": keypoints}, f, indent=4)
+        return "No keypoints detected.", None, None
 
     keypoints = np.array(keypoints[0])[:, :2]
 
-    NOSE, L_SHOULDER, R_SHOULDER = 0, 5, 6
-    L_WRIST = 9
-    L_HIP, R_HIP = 11, 12
-    L_KNEE, R_KNEE = 13, 14
-    L_ANKLE, R_ANKLE = 15, 16
-
-    skeleton = [(L_SHOULDER, R_SHOULDER), (L_SHOULDER, L_HIP), (R_SHOULDER, R_HIP), (L_HIP, R_HIP)]
-
-    # Draw keypoints and skeleton for visualization
+    # Draw keypoints and skeleton
+    skeleton = [(5, 6), (5, 11), (6, 12), (11, 12)]
     for x, y in keypoints:
-        cv2.circle(image_np, (int(x), int(y)), 5, (0, 255, 0), -1)
+        cv2.circle(image_cv, (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_np, (x1, y1), (x2, y2), (255, 0, 0), 2)
+        cv2.line(image_cv, (x1, y1), (x2, y2), (255, 0, 0), 2)
 
-    def get_distance(p1, p2):
-        return np.linalg.norm(np.array(p1) - np.array(p2))
+    # Body part indices from COCO keypoints
+    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
 
+    # Calculate pixel height based on nose to midpoint of ankles (rough estimate)
     ankle_mid = ((keypoints[L_ANKLE] + keypoints[R_ANKLE]) / 2).tolist()
     pixel_height = get_distance(keypoints[NOSE], ankle_mid)
-    estimated_full_pixel_height = pixel_height / 0.87
+    estimated_full_pixel_height = pixel_height / 0.87  # Adjust factor to estimate full height if partially visible
     pixels_per_cm = estimated_full_pixel_height / user_height_cm
 
+    # Measurements in 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
-
     pelvis = ((keypoints[L_HIP] + keypoints[R_HIP]) / 2).tolist()
-    neck = ((keypoints[L_SHOULDER] + keypoints[R_SHOULDER]) / 2).tolist()
-    torso_length_cm = get_distance(neck, pelvis) / pixels_per_cm
-
+    neck_point = ((keypoints[L_SHOULDER] + keypoints[R_SHOULDER]) / 2).tolist()
+    torso_length_cm = get_distance(neck_point, pelvis) / pixels_per_cm
     arm_length_cm = get_distance(keypoints[L_SHOULDER], keypoints[L_WRIST]) / pixels_per_cm
 
-    knee_mid = ((keypoints[L_KNEE] + keypoints[R_KNEE]) / 2).tolist()
-    neck_to_knee_cm = get_distance(neck, knee_mid) / pixels_per_cm
+    # Additional measurements
 
-    waist_circumference = np.pi * waist_width_cm
-    hip_circumference = waist_circumference / 0.75
+    # Neck circumference: approximate neck width * pi
+    neck_left = keypoints[L_SHOULDER]
+    neck_right = keypoints[R_SHOULDER]
+    neck_width_cm = get_distance(neck_left, neck_right) / pixels_per_cm
+    neck_circumference_cm = neck_width_cm * np.pi
+
+    # Calf circumference: approximate calf width at ankle * pi
+    ankle_left = keypoints[L_ANKLE]
+    ankle_right = keypoints[R_ANKLE]
+    calf_width_cm = get_distance(ankle_left, ankle_right) / pixels_per_cm
+    calf_circumference_cm = calf_width_cm * np.pi
+
+    # Waist and Hip circumferences
+    waist_circumference_cm = waist_width_cm * np.pi
+    hip_circumference_cm = waist_circumference_cm / 0.75  # Approximate relation
 
     measurements = {
+        "Waist Circumference (cm)": round(waist_circumference_cm, 2),
+        "Hip Circumference (cm)": round(hip_circumference_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)
+        "Full Arm Length (Shoulder to Wrist, cm)": round(arm_length_cm, 2),
+        "Neck Circumference (cm)": round(neck_circumference_cm, 2),
+        "Calf Circumference (cm)": round(calf_circumference_cm, 2),
+    }
+
+    return measurements, cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB), keypoints.tolist()
+
+
+def save_to_database(measurements, image, user_height_cm, user_id):
+    if not user_id:
+        return "❌ user_id missing from URL."
+    if measurements is None or image is None:
+        return "⚠️ No data to save."
+
+    buffered = BytesIO()
+    pil_image = Image.fromarray(image)
+    pil_image.save(buffered, format="JPEG")
+    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+    payload = {
+        "imageBase64": img_str,
+        "heightCm": user_height_cm,
+        "waistCircumferenceCm": measurements.get("Waist Circumference (cm)"),
+        "hipCircumferenceCm": measurements.get("Hip Circumference (cm)"),
+        "shoulderWidthCm": measurements.get("Shoulder Width (cm)"),
+        "torsoLengthCm": measurements.get("Torso Length (Neck to Pelvis, cm)"),
+        "fullArmLengthCm": measurements.get("Full Arm Length (Shoulder to Wrist, cm)"),
+        "neckCircumferenceCm": measurements.get("Neck Circumference (cm)"),
+        "calfCircumferenceCm": measurements.get("Calf Circumference (cm)"),
     }
 
-    return measurements, np.array(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB))
+    try:
+        response = requests.post(
+            f"https://7da2-2409-4042-6e81-1806-de6-b8e5-836c-6b95.ngrok-free.app/upload/{user_id}",
+            json=payload
+        )
+        if response.status_code == 201:
+            return "✅ Measurements and image saved to database!"
+        else:
+            return f"❌ Failed: {response.status_code} - {response.text}"
+    except Exception as e:
+        return f"⚠️ Error during save: {str(e)}"
+
 
-# Gradio Interface
+# === Gradio App ===
 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.")
+    gr.Markdown("# 📏 AI-Powered Body Measurement Tool")
+    user_id_state = gr.State()
+
+    @demo.load(inputs=None, outputs=[user_id_state])
+    def load_user_id(request: gr.Request):
+        return request.query_params.get("user_id", "")
 
     with gr.Row():
         with gr.Column():
-            image_input = gr.Image(type="pil", label="📸 Upload Image")
-            user_id_input = gr.Textbox(label="🆔 User ID", placeholder="Enter your user ID")
-            submit_btn = gr.Button("🔍 Generate Measurements")
+            image_input = gr.Image(label="Upload Your Full Body Image", type="pil")
+            height_input = gr.Number(label="Your Real Height (in cm)")
+            process_button = gr.Button("📐 Extract Measurements")
+
         with gr.Column():
-            height_input = gr.Number(label="📏 Your Height (cm)", value=170)
-            measurement_output = gr.JSON(label="📐 Estimated Measurements")
-            save_status = gr.Textbox(label="��� Save Status")
-
-    def process_and_save(image, height, user_id):
-        measurements, img_with_keypoints = process_image(image, height)
-        if isinstance(measurements, str):  # Means an error message
-            return measurements, None, ""
-        save_msg = save_to_database(measurements, np.array(img_with_keypoints), height, user_id)
-        return measurements, img_with_keypoints, save_msg
-
-    submit_btn.click(
-        fn=process_and_save,
-        inputs=[image_input, height_input, user_id_input],
-        outputs=[measurement_output, image_input, save_status]
+            output_image = gr.Image(label="Detected Keypoints")
+            measurement_output = gr.JSON(label="Body Measurements")
+
+    with gr.Row():
+        save_button = gr.Button("💾 Save to Backend")
+        save_status = gr.Textbox(label="Status", interactive=False)
+
+    # Store results for save
+    processed_img = gr.State()
+    processed_data = gr.State()
+
+    process_button.click(
+        fn=process_image,
+        inputs=[image_input, height_input],
+        outputs=[measurement_output, output_image, processed_data]
+    ).then(
+        fn=lambda img: img,
+        inputs=[output_image],
+        outputs=processed_img
     )
 
-demo.launch()
+    save_button.click(
+        fn=save_to_database,
+        inputs=[measurement_output, processed_img, height_input, user_id_state],
+        outputs=save_status
+    )
+
+if __name__ == "__main__":
+    demo.launch()
+