fix(fasterrcnn): handle imgs with 0 chunk

gnes-ai · Jul 18, 2019 · 92f36c3 · 92f36c3
1 parent 2da1942
commit 92f36c3
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/gnes/preprocessor/image/segmentation.py b/gnes/preprocessor/image/segmentation.py
@@ -42,6 +42,14 @@ def apply(self, doc: 'gnes_pb2.Document'):
             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.
+
             for ci, ele in enumerate(zip(chunks, weight)):
                 c = doc.chunks.add()
                 c.doc_id = doc.doc_id