Update app.py

app.py

@@ -4,6 +4,7 @@ import cv2
 import numpy as np
 import os
 import requests
+import torch
 
 # Ensure the model file is in the correct location
 model_path = "yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt"
@@ -14,18 +15,25 @@ if not os.path.exists(model_path):
     with open(model_path, "wb") as f:
         f.write(response.content)
 
-# Load the document segmentation model
-docseg_model = YOLO(model_path)
+# Load the document segmentation model on CPU
+device = torch.device('cpu')
+docseg_model = YOLO(model_path).to(device)
 
 def process_image(image):
     # Convert image to the format YOLO model expects
     image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-    results = docseg_model(source=image, save=False, show_labels=True, show_conf=True, show_boxes=True)
+    results = docseg_model(image)
 
     # Extract annotated image from results
     annotated_img = results[0].plot()
+    annotated_img = cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB)
 
-    return annotated_img, results[0].boxes
+    # Prepare detected areas and labels as text output
+    detected_areas_labels = "\n".join(
+        [f"{box.label}: {box.conf:.2f}" for box in results[0].boxes]
+    )
+
+    return annotated_img, detected_areas_labels
 
 # Define the Gradio interface
 interface = gr.Interface(