logo
Browse Source

Support cuda

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
Jael Gu 3 years ago
parent
commit
621ee3b17c
  1. 4
      detectron.py

4
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)

Loading…
Cancel
Save