wrap yolo pred with gpu

app.py

@@ -1,6 +1,7 @@
 # [DEBUG]
 #from osgeo import gdal
 
+
 import gradio as gr
 import logging
 from inference_tab import get_inference_widgets, run_inference

inference_tab/inference_logic.py

@@ -1,3 +1,4 @@
+import spaces
 import numpy as np
 from ultralytics import YOLO
 import os
@@ -52,7 +53,10 @@ def run_inference(image_path, gcp_path, city_name, score_th):
         else:
             yield msg, None
 
-    
+
+@spaces.GPU
+def yolo_pred(model,tile,tile_size,conf_threshold,verbose=False):
+    return model.predict(source=tile, imgsz=tile_size, conf=conf_threshold, verbose=verbose)    
 
 def getBBoxes(image_path, tile_size=256, overlap=0.3, confidence_threshold=0.25):
     yield f"DEBUG: Received image_path: {image_path}"
@@ -82,7 +86,7 @@ def getBBoxes(image_path, tile_size=256, overlap=0.3, confidence_threshold=0.25)
             if tile.shape[0] < tile_size or tile.shape[1] < tile_size:
                 continue
 
-            results = model.predict(source=tile, imgsz=tile_size, conf=confidence_threshold, verbose=False)
+            results = yolo_pred(model,tile,tile_size,confidence_threshold,verbose=False)
 
             for result in results:
                 boxes = result.boxes.xyxy.cpu().numpy()