Browse Source
Support cuda
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
1 changed files with
1 additions and
1 deletions
-
detectron.py
|
|
@ -54,7 +54,7 @@ class Detectron2(NNOperator): |
|
|
# Detectron2 uses BGR-formatted images |
|
|
# Detectron2 uses BGR-formatted images |
|
|
res = self._predictor(image[:,:,::-1]) |
|
|
res = self._predictor(image[:,:,::-1]) |
|
|
res = res['instances'] |
|
|
res = res['instances'] |
|
|
boxes = res.get('pred_boxes').tensor.numpy() |
|
|
|
|
|
|
|
|
boxes = res.get('pred_boxes').tensor.cpu().numpy() |
|
|
classes = res.get('pred_classes').numpy() |
|
|
classes = res.get('pred_classes').numpy() |
|
|
scores = res.get('scores').numpy() |
|
|
scores = res.get('scores').numpy() |
|
|
|
|
|
|
|
|
|
