Shikun Liu commited on
Commit
b0c2018
1 Parent(s): 58d758d

Parallelize expert label generation (#2)

Browse files

- parallelize expert label generation (e69a62beaadef055482ef76538964c81f8508323)

Files changed (1) hide show
  1. prismer_model.py +17 -8
prismer_model.py CHANGED
@@ -1,5 +1,6 @@
1
  from __future__ import annotations
2
 
 
3
  import os
4
  import pathlib
5
  import shlex
@@ -42,6 +43,18 @@ def build_deformable_conv() -> None:
42
  subprocess.run(shlex.split('sh make.sh'), cwd='prismer/experts/segmentation/mask2former/modeling/pixel_decoder/ops')
43
 
44
 
 
 
 
 
 
 
 
 
 
 
 
 
45
  def run_experts(image_path: str) -> tuple[str | None, ...]:
46
  helper_dir = submodule_dir / 'helpers'
47
  shutil.rmtree(helper_dir, ignore_errors=True)
@@ -50,14 +63,10 @@ def run_experts(image_path: str) -> tuple[str | None, ...]:
50
  out_path = image_dir / 'image.jpg'
51
  cv2.imwrite(out_path.as_posix(), cv2.imread(image_path))
52
 
53
- expert_names = ['depth', 'edge', 'normal', 'objdet', 'ocrdet', 'segmentation']
54
- for expert_name in expert_names:
55
- env = os.environ.copy()
56
- if 'PYTHONPATH' in env:
57
- env['PYTHONPATH'] = f'{submodule_dir.as_posix()}:{env["PYTHONPATH"]}'
58
- else:
59
- env['PYTHONPATH'] = submodule_dir.as_posix()
60
- subprocess.run(shlex.split(f'python experts/generate_{expert_name}.py'), cwd='prismer', env=env, check=True)
61
 
62
  keys = ['depth', 'edge', 'normal', 'seg_coco', 'obj_detection', 'ocr_detection']
63
  results = [pathlib.Path('prismer/helpers/labels') / key / 'helpers/images/image.png' for key in keys]
1
  from __future__ import annotations
2
 
3
+ import concurrent.futures
4
  import os
5
  import pathlib
6
  import shlex
43
  subprocess.run(shlex.split('sh make.sh'), cwd='prismer/experts/segmentation/mask2former/modeling/pixel_decoder/ops')
44
 
45
 
46
+ def run_expert(expert_name: str):
47
+ env = os.environ.copy()
48
+ if 'PYTHONPATH' in env:
49
+ env['PYTHONPATH'] = f'{submodule_dir.as_posix()}:{env["PYTHONPATH"]}'
50
+ else:
51
+ env['PYTHONPATH'] = submodule_dir.as_posix()
52
+ subprocess.run(shlex.split(f'python experts/generate_{expert_name}.py'),
53
+ cwd='prismer',
54
+ env=env,
55
+ check=True)
56
+
57
+
58
  def run_experts(image_path: str) -> tuple[str | None, ...]:
59
  helper_dir = submodule_dir / 'helpers'
60
  shutil.rmtree(helper_dir, ignore_errors=True)
63
  out_path = image_dir / 'image.jpg'
64
  cv2.imwrite(out_path.as_posix(), cv2.imread(image_path))
65
 
66
+ expert_names = ['edge', 'normal', 'objdet', 'ocrdet', 'segmentation']
67
+ run_expert('depth')
68
+ with concurrent.futures.ProcessPoolExecutor() as executor:
69
+ executor.map(run_expert, expert_names)
 
 
 
 
70
 
71
  keys = ['depth', 'edge', 'normal', 'seg_coco', 'obj_detection', 'ocr_detection']
72
  results = [pathlib.Path('prismer/helpers/labels') / key / 'helpers/images/image.png' for key in keys]