Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from ultralytics import YOLO
 import spaces
+import torch
 
 # Load pre-trained YOLOv8 model
 model = YOLO("yolov8x-doclaynet-epoch64-imgsz640-initiallr1e-4-finallr1e-5.pt")
@@ -8,8 +9,7 @@ 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) 
+@spaces.GPU(duration=60)
 def process_image(image):
     try:
         # Process the image
@@ -19,10 +19,10 @@ def process_image(image):
         # 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"{class_names[int(box.cls)].upper()}: {box.conf:.2f}" for box in result.boxes]
-        )
+        # Convert box.cls tensor to float before formatting
+        detected_areas_labels = "\n".join([
+            f"{class_names[int(box.cls.item())].upper()}: {box.conf:.2f}" for box in result.boxes
+        ])
 
         return annotated_image, detected_areas_labels
     except Exception as e: