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

gnes-ai · Jul 19, 2019 · 9f6c052 · 9f6c052
1 parent 3febe99
commit 9f6c052
Showing 1 changed file with 11 additions and 8 deletions.
diff --git a/gnes/preprocessor/image/segmentation.py b/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!')