Fix visualize script.

Author: IRONICBo

tools/visualize.py

@@ -4,42 +4,34 @@
 import os
 import os.path as osp
 import time
-from functools import partial
 
-from tqdm import tqdm
-import numpy as np
 import mmcv
 import mmengine
+import numpy as np
+from tqdm import tqdm
 
 from mmdeploy.apis import visualize_model
-from mmdeploy.utils import (Backend, get_backend, get_root_logger,
-                            load_config)
+from mmdeploy.utils import Backend, get_backend, get_root_logger, load_config
 
 
 def parse_args():
-    parser = argparse.ArgumentParser(description='Model inference visualization.')
+    parser = argparse.ArgumentParser(
+        description='Model inference visualization.')
     parser.add_argument('--deploy-cfg', help='deploy config path')
     parser.add_argument('--model-cfg', help='model config path')
-    parser.add_argument('--deploy-path', 
-        type=str,
-        nargs='+',
-        help='deploy model path')
     parser.add_argument(
-        '--checkpoint', 
-        default=None, 
-        help='model checkpoint path')
+        '--deploy-path', type=str, nargs='+', help='deploy model path')
+    parser.add_argument(
+        '--checkpoint', default=None, help='model checkpoint path')
     parser.add_argument(
         '--test-img',
         default=None,
         type=str,
         nargs='+',
         help='image used to test model')
     parser.add_argument(
-        '--save-dir',
-        default=None,
-        help='the dir to save inference results')
-    parser.add_argument(
-        '--device', help='device to run model', default='cpu')
+        '--save-dir', default=None, help='the dir to save inference results')
+    parser.add_argument('--device', help='device to run model', default='cpu')
     parser.add_argument(
         '--log-level',
         help='set log level',
@@ -65,61 +57,48 @@ def main():
     deploy_model_path = args.deploy_path
     if not isinstance(deploy_model_path, list):
         deploy_model_path = [deploy_model_path]
-    
+
     # load deploy_cfg
     deploy_cfg = load_config(deploy_cfg_path)[0]
-    
+
     # create save_dir or generate default save_dir
     save_dir = args.save_dir
     if save_dir:
         # generate default dir
         current_time = time.localtime()
-        save_dir = osp.join(
-            os.getcwd(), time.strftime("%Y_%m_%d_%H_%M_%S", current_time)
-        )
+        save_dir = osp.join(os.getcwd(),
+                            time.strftime('%Y_%m_%d_%H_%M_%S', current_time))
     mmengine.mkdir_or_exist(save_dir)
 
     # get backend info
     backend = get_backend(deploy_cfg)
     extra = dict()
     if backend == Backend.SNPE:
         extra['uri'] = args.uri
-    
+
     # iterate single_img
     for single_img in tqdm(args.test_img):
         filename = osp.basename(single_img)
-        output_file=osp.join(save_dir, filename)
-        visualize_model(
-            model_cfg_path,
-            deploy_cfg_path, 
-            deploy_model_path, 
-            single_img,
-            args.device,
-            backend,
-            output_file,
-            False,
-            **extra)
-                
+        output_file = osp.join(save_dir, filename)
+        visualize_model(model_cfg_path, deploy_cfg_path, deploy_model_path,
+                        single_img, args.device, backend, output_file, False,
+                        **extra)
+
         if checkpoint_path:
             pytorch_output_file = osp.join(save_dir, 'pytorch_out.jpg')
-            visualize_model(
-                model_cfg_path, 
-                deploy_cfg_path, 
-                [checkpoint_path],
-                single_img, 
-                args.device,
-                Backend.PYTORCH, 
-                pytorch_output_file, 
-                False)
-            
+            visualize_model(model_cfg_path, deploy_cfg_path, [checkpoint_path],
+                            single_img, args.device, Backend.PYTORCH,
+                            pytorch_output_file, False)
+
             # concat pytorch result and backend result
             backend_result = mmcv.imread(output_file)
             pytorch_result = mmcv.imread(pytorch_output_file)
             result = np.concatenate((backend_result, pytorch_result), axis=1)
             mmcv.imwrite(result, output_file)
-            
+
             # remove temp pytorch result
             os.remove(osp.join(save_dir, pytorch_output_file))
 
+
 if __name__ == '__main__':
-    main()
+    main()