Parallelize expert label generation

prismer_model.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import concurrent.futures
 import os
 import pathlib
 import shlex
@@ -42,6 +43,18 @@ def build_deformable_conv() -> None:
     subprocess.run(shlex.split('sh make.sh'), cwd='prismer/experts/segmentation/mask2former/modeling/pixel_decoder/ops')
 
 
+def run_expert(expert_name: str):
+    env = os.environ.copy()
+    if 'PYTHONPATH' in env:
+        env['PYTHONPATH'] = f'{submodule_dir.as_posix()}:{env["PYTHONPATH"]}'
+    else:
+        env['PYTHONPATH'] = submodule_dir.as_posix()
+    subprocess.run(shlex.split(f'python experts/generate_{expert_name}.py'),
+                   cwd='prismer',
+                   env=env,
+                   check=True)
+
+
 def run_experts(image_path: str) -> tuple[str | None, ...]:
     helper_dir = submodule_dir / 'helpers'
     shutil.rmtree(helper_dir, ignore_errors=True)
@@ -50,14 +63,10 @@ def run_experts(image_path: str) -> tuple[str | None, ...]:
     out_path = image_dir / 'image.jpg'
     cv2.imwrite(out_path.as_posix(), cv2.imread(image_path))
 
-    expert_names = ['depth', 'edge', 'normal', 'objdet', 'ocrdet', 'segmentation']
-    for expert_name in expert_names:
-        env = os.environ.copy()
-        if 'PYTHONPATH' in env:
-            env['PYTHONPATH'] = f'{submodule_dir.as_posix()}:{env["PYTHONPATH"]}'
-        else:
-            env['PYTHONPATH'] = submodule_dir.as_posix()
-        subprocess.run(shlex.split(f'python experts/generate_{expert_name}.py'), cwd='prismer', env=env, check=True)
+    expert_names = ['edge', 'normal', 'objdet', 'ocrdet', 'segmentation']
+    run_expert('depth')
+    with concurrent.futures.ProcessPoolExecutor() as executor:
+        executor.map(run_expert, expert_names)
 
     keys = ['depth', 'edge', 'normal', 'seg_coco', 'obj_detection', 'ocr_detection']
     results = [pathlib.Path('prismer/helpers/labels') / key / 'helpers/images/image.png' for key in keys]