add show

Author: wyh

synthesis/config.py

@@ -7,9 +7,13 @@
 config.datasets.CocoCVContoursDataset = AttrDict()
 config.datasets.CocoCVContoursDataset.auto_detect_subdir_with_basenum = 0
 
-config.sample_path_prefix = 'your sample path prefix'
-config.target_path = 'your json file path'
-config.output_dir = 'your output dir'
+sample_path_prefix_list = ['your sample path prefix list']
+target_path_list = ['your json file path list']
+config.output_label_dir = 'your output label dir'
+config.output_image_dir = 'your output image dir'
+config.show = True
 
-config.test_dataset = CocoCVContoursDataset(config, config.sample_path_prefix, config.target_path)
-config.test_loader = torch.utils.data.DataLoader(config.test_dataset, batch_size=1, pin_memory=False)
+config.test_loader_list = []
+for i in range(len(sample_path_prefix_list)):
+    test_dataset = CocoCVContoursDataset(config, sample_path_prefix_list[i], target_path_list[i])
+    config.test_loader_list.append(torch.utils.data.DataLoader(test_dataset, batch_size=1, pin_memory=False))

synthesis/synthesis.py

@@ -1,18 +1,41 @@
 import os
+import numpy as np
+import cv2
+from deepvac import LOG
 
-def synthesis(deepvac_config):
-    if not os.path.exists(deepvac_config.output_dir):
-        os.makedirs(deepvac_config.output_dir)
-    for i, (img, bbox_info, file_path) in enumerate(deepvac_config.test_loader):
-        file_path = file_path[0]
-        file_name = file_path.split('/')[-1]
-        txt_name = 'gt_{}'.format(file_name.replace('jpg', 'txt'))
-        fw = open(os.path.join(deepvac_config.output_dir, txt_name), 'w')
-        for s in bbox_info:
-            fw.write('{}\n'.format(s[0]))
-        fw.close()
+class Synthesis(object):
+    def __init__(self, deepvac_config):
+        self.config = deepvac_config
+        self.auditConfig()
+    
+    def auditConfig(self):
+        if not os.path.exists(self.config.output_label_dir):
+            os.makedirs(self.config.output_label_dir)
+        
+        if self.config.show and not os.path.exists(self.config.output_image_dir):
+            os.makedirs(self.config.output_image_dir)
+    
+    def synthesis(self):
+        for i, (img, bbox_info, file_path, show_img) in enumerate(self.test_loader):
+            file_path = file_path[0]
+            file_name = file_path.split('/')[-1]
+            txt_name = 'gt_{}'.format(file_name.replace('jpg', 'txt'))
+            fw = open(os.path.join(self.config.output_label_dir, txt_name), 'w')
+            for s in bbox_info:
+                fw.write('{}\n'.format(s[0]))
+            fw.close()
+
+            if self.config.show:
+                cv2.imwrite(os.path.join(self.config.output_image_dir, 'bbox_{}'.format(file_name)), show_img.numpy()[0])
+
+    def __call__(self):
+        for i, test_loader in enumerate(self.config.test_loader_list):
+            LOG.logI('Process dataset {}'.format(i+1))
+            self.test_loader = test_loader
+            self.synthesis()
 
 if __name__ == "__main__":
     from config import config as deepvac_config
 
-    synthesis(deepvac_config)
+    synthesis = Synthesis(deepvac_config)
+    synthesis()