Update app.py

app.py

@@ -4,10 +4,22 @@ import numpy as np
 import os
 import tempfile
 from ultralytics import YOLO
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Load the Latex2Layout model
 model_path = "latex2layout_object_detection_yolov8.pt"
-model = YOLO(model_path)
+try:
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(f"Model file not found: {model_path}")
+    model = YOLO(model_path)
+    logger.info("Model loaded successfully")
+except Exception as e:
+    logger.error(f"Error loading model: {str(e)}")
+    raise
 
 def detect_and_visualize(image):
     """
@@ -20,45 +32,56 @@ def detect_and_visualize(image):
         annotated_image: Image with detection boxes
         layout_annotations: Annotations in YOLO format
     """
-    if image is None:
-        return None, "Error: No image uploaded."
-    
-    # Run detection using the Latex2Layout model
-    results = model(image)
-    result = results[0]
-    
-    # Create a copy of the image for visualization
-    annotated_image = image.copy()
-    layout_annotations = []
-    
-    # Get image dimensions
-    img_height, img_width = image.shape[:2]
-    
-    # Draw detection results
-    for box in result.boxes:
-        x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
-        conf = float(box.conf[0])
-        cls_id = int(box.cls[0])
-        cls_name = result.names[cls_id]
+    try:
+        if image is None:
+            return None, "Error: No image uploaded."
         
-        # Generate a color for each class
-        color = tuple(np.random.randint(0, 255, 3).tolist())
+        # Validate image format and dimensions
+        if not isinstance(image, np.ndarray):
+            return None, "Error: Invalid image format."
         
-        # Draw bounding box and label
-        cv2.rectangle(annotated_image, (x1, y1), (x2, y2), color, 2)
-        label = f'{cls_name} {conf:.2f}'
-        (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
-        cv2.rectangle(annotated_image, (x1, y1-label_height-5), (x1+label_width, y1), color, -1)
-        cv2.putText(annotated_image, label, (x1, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
+        if image.size == 0:
+            return None, "Error: Empty image."
         
-        # Convert to YOLO format (normalized)
-        x_center = (x1 + x2) / (2 * img_width)
-        y_center = (y1 + y2) / (2 * img_height)
-        width = (x2 - x1) / img_width
-        height = (y2 - y1) / img_height
-        layout_annotations.append(f"{cls_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
-    
-    return annotated_image, "\n".join(layout_annotations)
+        # Run detection using the Latex2Layout model
+        results = model(image)
+        result = results[0]
+        
+        # Create a copy of the image for visualization
+        annotated_image = image.copy()
+        layout_annotations = []
+        
+        # Get image dimensions
+        img_height, img_width = image.shape[:2]
+        
+        # Draw detection results
+        for box in result.boxes:
+            x1, y1, x2, y2 = map(int, box.xyxy[0].cpu().numpy())
+            conf = float(box.conf[0])
+            cls_id = int(box.cls[0])
+            cls_name = result.names[cls_id]
+            
+            # Generate a color for each class
+            color = tuple(np.random.randint(0, 255, 3).tolist())
+            
+            # Draw bounding box and label
+            cv2.rectangle(annotated_image, (x1, y1), (x2, y2), color, 2)
+            label = f'{cls_name} {conf:.2f}'
+            (label_width, label_height), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
+            cv2.rectangle(annotated_image, (x1, y1-label_height-5), (x1+label_width, y1), color, -1)
+            cv2.putText(annotated_image, label, (x1, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
+            
+            # Convert to YOLO format (normalized)
+            x_center = (x1 + x2) / (2 * img_width)
+            y_center = (y1 + y2) / (2 * img_height)
+            width = (x2 - x1) / img_width
+            height = (y2 - y1) / img_height
+            layout_annotations.append(f"{cls_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
+        
+        return annotated_image, "\n".join(layout_annotations)
+    except Exception as e:
+        logger.error(f"Error in detect_and_visualize: {str(e)}")
+        return None, f"Error during detection: {str(e)}"
 
 def save_layout_annotations(layout_annotations_str):
     """
@@ -70,13 +93,34 @@ def save_layout_annotations(layout_annotations_str):
     Returns:
         file_path: Path to the saved annotation file
     """
-    if not layout_annotations_str:
+    try:
+        if not layout_annotations_str:
+            return None
+        
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".txt")
+        with open(temp_file.name, "w") as f:
+            f.write(layout_annotations_str)
+        return temp_file.name
+    except Exception as e:
+        logger.error(f"Error in save_layout_annotations: {str(e)}")
         return None
+
+def load_example_image():
+    """
+    Load an example image for demonstration.
     
-    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".txt")
-    with open(temp_file.name, "w") as f:
-        f.write(layout_annotations_str)
-    return temp_file.name
+    Returns:
+        image: The loaded example image or None if loading fails
+    """
+    try:
+        example_path = "example_image.jpg"
+        if not os.path.exists(example_path):
+            logger.error(f"Example image not found: {example_path}")
+            return None
+        return cv2.imread(example_path)
+    except Exception as e:
+        logger.error(f"Error loading example image: {str(e)}")
+        return None
 
 # Custom CSS for styling
 custom_css = """
@@ -144,7 +188,7 @@ with gr.Blocks(
     # Example image button (optional)
     with gr.Row():
         gr.Button("Load Example Image").click(
-            fn=lambda: cv2.imread("example_image.jpg"),
+            fn=load_example_image,
             outputs=input_image
         )
     
@@ -153,12 +197,7 @@ with gr.Blocks(
         fn=detect_and_visualize,
         inputs=input_image,
         outputs=[output_image, layout_annotations],
-        _js="() => { document.querySelector('.button-primary').innerText = 'Processing...'; }",
         show_progress=True
-    ).then(
-        fn=lambda: gr.update(value="Start Detection"),
-        outputs=detect_btn,
-        _js="() => { document.querySelector('.button-primary').innerText = 'Start Detection'; }"
     )
     
     download_btn.click(