Spaces:
Running
Running
Martin Tomov
commited on
Commit
β’
027d733
1
Parent(s):
f44cc10
Update detectron_utils.py
Browse files- detectron_utils.py +6 -3
detectron_utils.py
CHANGED
@@ -1,4 +1,3 @@
|
|
1 |
-
import spaces
|
2 |
import torch
|
3 |
import numpy as np
|
4 |
import cv2
|
@@ -8,6 +7,9 @@ REPO_ID = "idml/Detectron2-FasterRCNN_InsectDetect"
|
|
8 |
FILENAME = "model.pth"
|
9 |
FILENAME_CONFIG = "config.yml"
|
10 |
|
|
|
|
|
|
|
11 |
import cv2
|
12 |
from detectron2.config import get_cfg
|
13 |
from detectron2.engine import DefaultPredictor
|
@@ -102,18 +104,19 @@ viz_classes = {'thing_classes': ['Acrididae',
|
|
102 |
|
103 |
|
104 |
|
105 |
-
@spaces.GPU
|
106 |
def detectron_process_image(image):
|
107 |
cfg = get_cfg()
|
108 |
|
|
|
109 |
cfg.merge_from_file(hf_hub_download(repo_id=REPO_ID, filename=FILENAME_CONFIG))
|
110 |
cfg.MODEL.WEIGHTS = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
|
111 |
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.2
|
112 |
-
cfg.MODEL.DEVICE
|
113 |
predictor = DefaultPredictor(cfg)
|
114 |
|
115 |
numpy_image = np.array(image)
|
116 |
|
|
|
117 |
im = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
|
118 |
|
119 |
v = Visualizer(im[:, :, ::-1],
|
|
|
|
|
1 |
import torch
|
2 |
import numpy as np
|
3 |
import cv2
|
|
|
7 |
FILENAME = "model.pth"
|
8 |
FILENAME_CONFIG = "config.yml"
|
9 |
|
10 |
+
|
11 |
+
# Ensure you have the model file
|
12 |
+
|
13 |
import cv2
|
14 |
from detectron2.config import get_cfg
|
15 |
from detectron2.engine import DefaultPredictor
|
|
|
104 |
|
105 |
|
106 |
|
|
|
107 |
def detectron_process_image(image):
|
108 |
cfg = get_cfg()
|
109 |
|
110 |
+
|
111 |
cfg.merge_from_file(hf_hub_download(repo_id=REPO_ID, filename=FILENAME_CONFIG))
|
112 |
cfg.MODEL.WEIGHTS = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
|
113 |
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.2
|
114 |
+
cfg.MODEL.DEVICE='cpu'
|
115 |
predictor = DefaultPredictor(cfg)
|
116 |
|
117 |
numpy_image = np.array(image)
|
118 |
|
119 |
+
|
120 |
im = cv2.cvtColor(numpy_image, cv2.COLOR_RGB2BGR)
|
121 |
|
122 |
v = Visualizer(im[:, :, ::-1],
|