Update app.py

app.py

@@ -17,7 +17,7 @@ def get_class_color(class_id):
     return CLASS_COLORS[class_id]
 
 # Class Names (Modify based on your dataset)
-CLASS_NAMES = {0: "Text Line", 1: "Heading", 2: "Signature"}  # Example labels
+CLASS_NAMES = {0: "Text", 1: "Non-Text"}  # Example labels
 
 def detect_text_lines(image):
     """Detects text lines with two different confidence and IoU thresholds."""
@@ -38,8 +38,8 @@ def detect_text_lines(image):
 
         # Run YOLO text detection with specific thresholds
         results = model.predict(image, conf=conf, iou=iou, device="cpu")
-        detected_boxes = results[0].boxes.xyxy.tolist()
-        class_ids = results[0].boxes.cls.tolist()
+        detected_boxes = results[0].boxes.xyxy.tolist() if hasattr(results[0].boxes, 'xyxy') else []
+        class_ids = results[0].boxes.cls.tolist() if hasattr(results[0].boxes, 'cls') else []
         detected_boxes = [list(map(int, box)) for box in detected_boxes]
 
         # Draw bounding boxes on the image
@@ -52,7 +52,7 @@ def detect_text_lines(image):
             font = ImageFont.load_default()  # Fallback in case font is missing
 
         for idx, (x1, y1, x2, y2) in enumerate(detected_boxes):
-            class_id = int(class_ids[idx])
+            class_id = int(class_ids[idx]) if idx < len(class_ids) else -1
             color = get_class_color(class_id)
             class_name = CLASS_NAMES.get(class_id, f"Class {class_id}")