From 669ac8f2f671b0ed186e2586859af07d8edd404b Mon Sep 17 00:00:00 2001
From: zilliz <zilliz@zillizs-MacBook-Pro-2.local>
Date: Tue, 26 Jul 2022 15:54:28 +0800
Subject: [PATCH] fix detecron2 install issue

---
 __init__.py  | 7 +++++++
 detectron.py | 2 +-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/__init__.py b/__init__.py
index a635800..f459787 100644
--- a/__init__.py
+++ b/__init__.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
 
 
diff --git a/detectron.py b/detectron.py
index 52e4c93..c3539ce 100644
--- a/detectron.py
+++ b/detectron.py
@@ -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()