Update app.py

app.py

@@ -15,21 +15,30 @@ app.secret_key = 'your_secret_key'  # Replace with a secure secret key
 #########################################
 
 # --- Roboflow Box Detection Model ---
-API_KEY = "wLjPoPYaLmrqCIOFA0RH"            # Replace with your actual API key
-PROJECT_ID = "base-model-box-r4suo-8lkk1-6dbqh"      # Replace with your Roboflow project ID
-VERSION_NUMBER = "2"  # Replace with your trained model version number
+API_KEY = "wLjPoPYaLmrqCIOFA0RH"           # Your Roboflow API key
+PROJECT_ID = "base-model-box-r4suo-8lkk1-6dbqh"  # Your Roboflow project ID
+VERSION_NUMBER = "2"                        # Your trained model version number
 
-rf = roboflow.Roboflow(api_key=API_KEY)
-workspace = rf.workspace()
-project = workspace.project(PROJECT_ID)
-version = project.version(VERSION_NUMBER)
-box_model = version.model  # This model is trained for detecting boxes
+try:
+    rf = roboflow.Roboflow(api_key=API_KEY)
+    workspace = rf.workspace()
+    project = workspace.project(PROJECT_ID)
+    version = project.version(VERSION_NUMBER)
+    box_model = version.model  # This model is trained for detecting boxes
+    print("Roboflow model loaded successfully.")
+except Exception as e:
+    print("Error initializing Roboflow model:", e)
+    box_model = None
 
 # --- YOLOv5 Pretrained Model for Persons & Cars ---
-# Using Ultralytics YOLOv5s (pretrained) from Torch Hub
-yolov5_model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
-# We'll filter YOLO detections to only include persons and cars.
-YOLO_FILTER_CLASSES = {"person", "car"}
+try:
+    yolov5_model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
+    # Filter YOLO detections to only include persons and cars.
+    YOLO_FILTER_CLASSES = {"person", "car"}
+    print("YOLOv5 model loaded successfully.")
+except Exception as e:
+    print("Error loading YOLOv5 model:", e)
+    yolov5_model = None
 
 #########################################
 # 2. Helper Functions
@@ -63,55 +72,68 @@ def custom_nms(preds, iou_threshold=0.3):
     return filtered_preds
 
 def process_image(image_path):
-    """
-    Process the uploaded image using both detection pipelines:
-      (a) Box detection via Roboflow (with measurement using an ArUco marker).
-      (b) YOLOv5 detection for persons and cars.
-    Returns the annotated image and a list of detection info dictionaries.
-    """
+    # Load image
     image = cv2.imread(image_path)
     if image is None:
+        print("DEBUG: cv2.imread failed to load the image from", image_path)
         return None, "Could not read the image."
+
     img_height, img_width = image.shape[:2]
-    
     detection_info = []  # List to hold all detection results for display
 
     # --- (a) Roboflow Box Detection & Measurement ---
-    results = box_model.predict(image_path, confidence=50, overlap=30).json()
-    predictions = results.get("predictions", [])
+    if box_model is None:
+        print("DEBUG: Roboflow model is not initialized.")
+        return None, "Roboflow model is not available."
+    try:
+        results = box_model.predict(image_path, confidence=50, overlap=30).json()
+        predictions = results.get("predictions", [])
+    except Exception as e:
+        print("DEBUG: Error during Roboflow prediction:", e)
+        return None, "Error during Roboflow prediction."
+
     processed_preds = []
     for prediction in predictions:
-        x, y, width, height = prediction["x"], prediction["y"], prediction["width"], prediction["height"]
-        x1 = int(round(x - width / 2))
-        y1 = int(round(y - height / 2))
-        x2 = int(round(x + width / 2))
-        y2 = int(round(y + height / 2))
-        # Clamp coordinates to image dimensions
-        x1 = max(0, min(x1, img_width - 1))
-        y1 = max(0, min(y1, img_height - 1))
-        x2 = max(0, min(x2, img_width - 1))
-        y2 = max(0, min(y2, img_height - 1))
-        processed_preds.append({
-            "box": (x1, y1, x2, y2),
-            "class": prediction["class"],
-            "confidence": prediction["confidence"]
-        })
+        try:
+            x, y, width, height = prediction["x"], prediction["y"], prediction["width"], prediction["height"]
+            x1 = int(round(x - width / 2))
+            y1 = int(round(y - height / 2))
+            x2 = int(round(x + width / 2))
+            y2 = int(round(y + height / 2))
+            # Clamp coordinates to image dimensions
+            x1 = max(0, min(x1, img_width - 1))
+            y1 = max(0, min(y1, img_height - 1))
+            x2 = max(0, min(x2, img_width - 1))
+            y2 = max(0, min(y2, img_height - 1))
+            processed_preds.append({
+                "box": (x1, y1, x2, y2),
+                "class": prediction["class"],
+                "confidence": prediction["confidence"]
+            })
+        except Exception as e:
+            print("DEBUG: Error processing a prediction:", e)
+            continue
+
     box_detections = custom_nms(processed_preds, iou_threshold=0.3)
-    
+
     # Detect ArUco marker for measurement (only applicable for boxes)
     marker_real_width_cm = 10.0  # The marker is 10cm x 10cm
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_6X6_250)
-    aruco_params = cv2.aruco.DetectorParameters()
-    corners, ids, _ = cv2.aruco.detectMarkers(gray, aruco_dict, parameters=aruco_params)
-    if ids is not None and len(corners) > 0:
-        marker_corners = corners[0].reshape((4, 2))
-        cv2.aruco.drawDetectedMarkers(image, corners, ids)
-        marker_width_pixels = np.linalg.norm(marker_corners[0] - marker_corners[1])
-        marker_height_pixels = np.linalg.norm(marker_corners[1] - marker_corners[2])
-        marker_pixel_size = (marker_width_pixels + marker_height_pixels) / 2.0
-        conversion_factor = marker_real_width_cm / marker_pixel_size
-    else:
+    try:
+        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+        aruco_dict = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_6X6_250)
+        aruco_params = cv2.aruco.DetectorParameters()
+        corners, ids, _ = cv2.aruco.detectMarkers(gray, aruco_dict, parameters=aruco_params)
+        if ids is not None and len(corners) > 0:
+            marker_corners = corners[0].reshape((4, 2))
+            cv2.aruco.drawDetectedMarkers(image, corners, ids)
+            marker_width_pixels = np.linalg.norm(marker_corners[0] - marker_corners[1])
+            marker_height_pixels = np.linalg.norm(marker_corners[1] - marker_corners[2])
+            marker_pixel_size = (marker_width_pixels + marker_height_pixels) / 2.0
+            conversion_factor = marker_real_width_cm / marker_pixel_size
+        else:
+            conversion_factor = None
+    except Exception as e:
+        print("DEBUG: Error during ArUco detection:", e)
         conversion_factor = None
 
     # Draw box detections and record measurement info (only for boxes)
@@ -144,32 +166,38 @@ def process_image(image_path):
         (text_width, text_height), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
         cv2.rectangle(image, (x1, y1 - text_height - baseline - 5), (x1 + text_width, y1 - 5), (0, 255, 0), -1)
         cv2.putText(image, text, (x1, y1 - 5 - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
-    
+
     # --- (b) YOLOv5 for Persons & Cars ---
-    # Convert image to RGB for YOLO (it expects RGB)
-    img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    yolo_results = yolov5_model(img_rgb)
-    df = yolo_results.pandas().xyxy[0]
-    for _, row in df.iterrows():
-        if row['name'] in YOLO_FILTER_CLASSES:
-            xmin = int(row['xmin'])
-            ymin = int(row['ymin'])
-            xmax = int(row['xmax'])
-            ymax = int(row['ymax'])
-            conf = row['confidence']
-            label = row['name']
-            cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (255, 0, 0), 2)
-            text = f"{label} ({conf:.2f})"
-            (text_width, text_height), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
-            cv2.rectangle(image, (xmin, ymin - text_height - baseline - 5), (xmin + text_width, ymin - 5), (255, 0, 0), -1)
-            cv2.putText(image, text, (xmin, ymin - 5 - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
-            detection_info.append({
-                "class": label,
-                "confidence": f"{conf:.2f}",
-                "width_cm": "N/A",
-                "height_cm": "N/A"
-            })
-    
+    if yolov5_model is None:
+        print("DEBUG: YOLOv5 model is not initialized.")
+    else:
+        try:
+            img_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            yolo_results = yolov5_model(img_rgb)
+            df = yolo_results.pandas().xyxy[0]
+            for _, row in df.iterrows():
+                if row['name'] in YOLO_FILTER_CLASSES:
+                    xmin = int(row['xmin'])
+                    ymin = int(row['ymin'])
+                    xmax = int(row['xmax'])
+                    ymax = int(row['ymax'])
+                    conf = row['confidence']
+                    label = row['name']
+                    cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (255, 0, 0), 2)
+                    text = f"{label} ({conf:.2f})"
+                    (text_width, text_height), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
+                    cv2.rectangle(image, (xmin, ymin - text_height - baseline - 5), (xmin + text_width, ymin - 5), (255, 0, 0), -1)
+                    cv2.putText(image, text, (xmin, ymin - 5 - baseline), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
+                    detection_info.append({
+                        "class": label,
+                        "confidence": f"{conf:.2f}",
+                        "width_cm": "N/A",
+                        "height_cm": "N/A"
+                    })
+        except Exception as e:
+            print("DEBUG: Error during YOLOv5 inference:", e)
+            return None, "Error during YOLOv5 inference."
+
     # --- Build Top Summary Text ---
     detection_counts = Counter(det["class"] for det in detection_info)
     if detection_counts:
@@ -177,7 +205,7 @@ def process_image(image_path):
         (info_width, info_height), info_baseline = cv2.getTextSize(top_text, cv2.FONT_HERSHEY_SIMPLEX, 1, 2)
         cv2.rectangle(image, (5, 5), (5 + info_width, 5 + info_height + info_baseline), (0, 255, 0), -1)
         cv2.putText(image, top_text, (5, 5 + info_height), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
-    
+
     return image, detection_info
 
 #########################################
@@ -197,14 +225,22 @@ def index():
             flash('No selected file')
             return redirect(request.url)
         upload_path = "uploaded.jpg"
-        file.save(upload_path)
+        try:
+            file.save(upload_path)
+        except Exception as e:
+            print("DEBUG: Error saving uploaded file:", e)
+            flash("Error saving uploaded file.")
+            return redirect(request.url)
         processed_image, detection_info = process_image(upload_path)
         if processed_image is None:
-            flash("Error processing image.")
+            flash("Error Processing Image: " + detection_info)
         else:
             retval, buffer = cv2.imencode('.jpg', processed_image)
             image_data = base64.b64encode(buffer).decode('utf-8')
-        os.remove(upload_path)
+        try:
+            os.remove(upload_path)
+        except Exception as e:
+            print("DEBUG: Error removing uploaded file:", e)
     return render_template_string('''
     <!doctype html>
     <html>
@@ -334,4 +370,5 @@ def index():
 #########################################
 
 if __name__ == '__main__':
+    # Ensure the app runs on 0.0.0.0 and port 7860 for Hugging Face Spaces.
     app.run(host="0.0.0.0", port=7860)