@spaces.GPU detectron_utils.py

detectron_utils.py

@@ -2,14 +2,12 @@ import torch
 import numpy as np
 import cv2
 from huggingface_hub import hf_hub_download
+import spaces
 
 REPO_ID = "idml/Detectron2-FasterRCNN_InsectDetect"
 FILENAME = "model.pth"
 FILENAME_CONFIG = "config.yml"
 
-
-# Ensure you have the model file
-
 import cv2
 from detectron2.config import get_cfg
 from detectron2.engine import DefaultPredictor
@@ -104,19 +102,18 @@ viz_classes = {'thing_classes': ['Acrididae',
 
 
 
+@spaces.GPU
 def detectron_process_image(image):
     cfg = get_cfg()
 
-
     cfg.merge_from_file(hf_hub_download(repo_id=REPO_ID, filename=FILENAME_CONFIG))
     cfg.MODEL.WEIGHTS = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
     cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.2
-    cfg.MODEL.DEVICE='cpu'
+    cfg.MODEL.DEVICE = 'cpu'
     predictor = DefaultPredictor(cfg)
 
     numpy_image = np.array(image)
 
-
     im = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
     
     v = Visualizer(im[:, :, ::-1],