Update app.py

app.py

@@ -1,28 +1,34 @@
 import gradio as gr
 from ultralytics import YOLO
-import spaces  # Import the `spaces` library
+import spaces
 
 # Load pre-trained YOLOv8 model
 model = YOLO("yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt")
 
+# Get class names from model
+class_names = model.names
+
 # Decorate the `process_image` function with `@spaces.GPU`
-@spaces.GPU(duration=60)  # Optional: Set the duration if needed
+@spaces.GPU(duration=60) 
 def process_image(image):
     try:
         # Process the image
         results = model(source=image, save=False, show_labels=True, show_conf=True, show_boxes=True)
-        result = results[0]
+        result = results[0]  # Get the first result
 
-        # Extract the annotated image and the labels/confidence scores
+        # Extract annotated image and labels with class names
         annotated_image = result.plot()
+
+        # Use cls attribute for labels and get class name from model
         detected_areas_labels = "\n".join(
-            [f"{box.label.upper()}: {box.conf:.2f}" for box in result.boxes]
+            [f"{class_names[int(box.cls)].upper()}: {box.conf:.2f}" for box in result.boxes]
         )
 
         return annotated_image, detected_areas_labels
     except Exception as e:
         return None, f"Error processing image: {e}"
 
+
 # Create the Gradio Interface
 with gr.Blocks() as demo:
     gr.Markdown("# Document Segmentation Demo (ZeroGPU)")