Spaces:
Sleeping
Sleeping
Martin Tomov
commited on
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],
|