zilliz
3 years ago
2 changed files with
8 additions and
1 deletions
-
__init__.py
-
detectron.py
|
|
|
@ -1,3 +1,10 @@ |
|
|
|
import os |
|
|
|
try: |
|
|
|
import detectron2 |
|
|
|
except ModuleNotFoundError: |
|
|
|
os.system("git clone https://github.com/facebookresearch/detectron2.git") |
|
|
|
os.system("python -m pip install -e detectron2") |
|
|
|
|
|
|
|
from .detectron import Detectron2 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@ -51,7 +51,7 @@ class Detectron2(NNOperator): |
|
|
|
def __call__(self, image: 'towhee._types.Image') -> Tuple[List]: |
|
|
|
|
|
|
|
# Detectron2 uses BGR-formatted images |
|
|
|
res = self._predictor(image.to_ndarray()[:,:,::-1]) |
|
|
|
res = self._predictor(image[:,:,::-1]) |
|
|
|
res = res['instances'] |
|
|
|
boxes = res.get('pred_boxes').tensor.numpy() |
|
|
|
classes = res.get('pred_classes').numpy() |
|
|
|
|
