refactor(inferencer): inference and visualize image one by one (#2)

ModelTC · Nov 8, 2021 · 0138d0c · 0138d0c
1 parent b0ac29a
commit 0138d0c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/eod/apis/inference.py b/eod/apis/inference.py
@@ -185,8 +185,8 @@ def predict(self):
             batch = self.fetch([filename])
             with torch.no_grad():
                 output = self.detector(batch)
-            output_list += self.map_back(output)
-        return output_list
+            output = self.map_back(output)
+            self.vis(output)
 
     def vis(self, outputs):
         for img_idx, output in enumerate(outputs):
@@ -198,7 +198,7 @@ def vis(self, outputs):
             filename = os.path.basename(output['image_info'][-1])
             if self.vis_type == 'plt':
                 filename = filename.rsplit('.', 1)[0]
-            logger.info('visualizing {}:{}'.format(img_idx, filename))
+            logger.info('visualizing {}'.format(filename))
 
             img_h, img_w = img.shape[:2]
             classes = boxes[:, -1].astype(np.int32)

diff --git a/eod/commands/inference.py b/eod/commands/inference.py
@@ -61,8 +61,7 @@ def main(args):
     infer_cfg['kwargs'] = infer_cfg.get('kwargs', {})
     cfg['runtime']['inferencer'] = infer_cfg
     inferencer = INFERENCER_REGISTRY.get(infer_cfg['type'])(cfg, **infer_cfg['kwargs'])
-    output_list = inferencer.predict()
-    inferencer.vis(output_list)
+    inferencer.predict()
 
 
 def _main(args):