Skip to content
Permalink
Browse files

fix(fasterrcnn): fix bug for gpu

  • Loading branch information...
jemmyshin committed Jul 17, 2019
1 parent f2667bb commit a13299132b776ccbb1826bf0f3a361e431a0f2ca
Showing with 2 additions and 1 deletion.
  1. +2 −1 gnes/preprocessor/image/segmentation.py
@@ -17,7 +17,6 @@ def __init__(self, model_name: str,
self.model_name = model_name
self.model_dir = model_dir
self.target_img_size = target_img_size
self.model_name = model_name
self._use_cuda = _use_cuda

def post_init(self):
@@ -37,6 +36,8 @@ def apply(self, doc: 'gnes_pb2.Document'):
if doc.raw_bytes:
original_image = Image.open(io.BytesIO(doc.raw_bytes))
image_tensor = self._torch_transform(original_image)
if self._use_cuda:
image_tensor = image_tensor.cuda()

seg_output = self._model([image_tensor])
chunks = seg_output[0]['boxes'].tolist()

0 comments on commit a132991

Please sign in to comment.
You can’t perform that action at this time.