Update app.py

app.py

@@ -6,8 +6,14 @@ import spaces
 
 # Load the trained model
 
-model_path = 'best.engine'  # Replace with the path to your trained .pt file
-model = YOLO(model_path,task="classify")
+# Load the YOLOv8 model
+pt_model = YOLO("best.pt")
+
+# Export the model to TensorRT format
+pt_model.export(format="engine")  # creates 'yolov8n.engine'
+
+# Load the exported TensorRT model
+model = YOLO("best.engine")
 
 # Define the class indices for figures and tables
 figure_class_index = 3  # class index for figures