Update app.py

app.py

@@ -132,16 +132,22 @@ import base64
 import requests
 from io import BytesIO
 
-# Function to process the image and calculate measurements
-def process_image(image, user_height_cm):
-    # Check if the user height is valid
-    if user_height_cm == 0 or user_height_cm is None:
-        return "User height cannot be zero or None. Please provide a valid height.", None
 
+
+# Load pre-trained Keypoint R-CNN 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"
+
+# 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 (using your predictor)
+    # Run keypoint detection
     outputs = predictor(image)
 
     # Extract keypoints
@@ -152,7 +158,6 @@ def process_image(image, user_height_cm):
         return "No keypoints detected.", None
 
     # Save keypoints to JSON
-    output_file = "keypoints.json"
     with open(output_file, "w") as f:
         json.dump({"keypoints": keypoints}, f, indent=4)
 
@@ -186,18 +191,10 @@ def process_image(image, user_height_cm):
     ankle_mid = ((keypoints[L_ANKLE] + keypoints[R_ANKLE]) / 2).tolist()
     pixel_height = get_distance(keypoints[NOSE], ankle_mid)
 
-    # Check if pixel height is non-zero to avoid divide by zero error
-    if pixel_height == 0:
-        return "Error in calculating pixel height. Please check the input image.", None
-
     # Estimated full body height (add approx head length)
     estimated_full_pixel_height = pixel_height / 0.87  # Since 87% = nose to ankle
     pixels_per_cm = estimated_full_pixel_height / user_height_cm
 
-    # Ensure pixels_per_cm is valid (non-zero)
-    if pixels_per_cm == 0:
-        return "Error in calculating pixels per cm. Ensure the user height is correct.", None
-
     # 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])
@@ -234,6 +231,7 @@ def process_image(image, user_height_cm):
     measurements = calculate_body_measurements(waist_circumference, hip_circumference, shoulder_width_cm, torso_length_cm, arm_length_cm)
 
     return measurements, cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+   
 
 # Save to database function
 def save_to_database(measurements, image, user_height_cm):