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