Update app.py

app.py

@@ -10,12 +10,12 @@ from PIL import Image
 model_path = 'best.pt'  # Replace with the path to your trained .pt file
 model = YOLO(model_path)
 
-# Define the class indices for figures and tables (adjust based on your model's classes)
+# Define the class indices for figures and tables
 figure_class_index = 3  # class index for figures
 table_class_index = 4   # class index for tables
 
 # Function to perform inference on an image and return bounding boxes for figures and tables
-def infer_image_and_get_boxes(image):
+def infer_image_and_get_boxes(image, confidence_threshold=0.6):
     # Convert the image from BGR to RGB
     image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
     
@@ -27,7 +27,8 @@ def infer_image_and_get_boxes(image):
     for result in results:
         for box in result.boxes:
             cls = int(box.cls[0])
-            if cls == figure_class_index or cls == table_class_index:
+            confidence = box.conf[0]
+            if (cls == figure_class_index or cls == table_class_index) and confidence > confidence_threshold:
                 x1, y1, x2, y2 = map(int, box.xyxy[0])
                 boxes.append((x1, y1, x2, y2))