Skip to content
Permalink
Browse files

feat(fasterrcnn): add the original image to chunk list

  • Loading branch information...
jemmyshin committed Jul 19, 2019
1 parent 3febe99 commit 9f6c0524f986d735d3ef831fe2021579dff575c3
Showing with 11 additions and 8 deletions.
  1. +11 −8 gnes/preprocessor/image/segmentation.py
@@ -40,22 +40,25 @@ def apply(self, doc: 'gnes_pb2.Document'):
image_tensor = image_tensor.cuda()

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

weight = seg_output[0]['scores'].tolist()
if len(chunks) == 0:
c = doc.chunks.add()
c.doc_id = doc.doc_id
c.blob.CopyFrom(array2blob(np.array(original_image.resize((self.target_img_size,
self.target_img_size)))))
c.offset_1d = 1
c.weight = 1.
length = len(list(filter(lambda x: x >= 0.5, weight)))
chunks = seg_output[0]['boxes'].tolist()[:length]
weight = weight[:length]

for ci, ele in enumerate(zip(chunks, weight)):
c = doc.chunks.add()
c.doc_id = doc.doc_id
c.blob.CopyFrom(array2blob(self._crop_image_reshape(original_image, ele[0])))
c.offset_1d = ci
c.weight = ele[1]

c = doc.chunks.add()
c.doc_id = doc.doc_id
c.blob.CopyFrom(array2blob(np.array(original_image.resize((self.target_img_size,
self.target_img_size)))))
c.offset_1d = len(chunks)
c.weight = 1.
else:
self.logger.error('bad document: "raw_bytes" is empty!')

0 comments on commit 9f6c052

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