mp_gen_res.py

from tqdm import tqdm
import subprocess
import sys
import os
from multiprocessing import Pool, cpu_count, Manager
from time import perf_counter

if len(sys.argv) != 3:
    print('usage: gen_res.py /path/to/image /path/to/output')

curpath = os.path.dirname(os.path.abspath(__file__))
path_to_image = sys.argv[1]


def process_single_img(path_to_image, filename, output_dict):
    img_path = os.path.join(path_to_image, filename)
    detection = '/home/CCTag/build/Linux-x86_64/detection'
    det_flags = ['-n', '3',           # three concentric circles
                 '-i', img_path]      # input image
    cmd = [detection, *det_flags, img_path, sys.argv[2]]
    with open(os.devnull, 'w') as devnull:
        det_result = \
            subprocess.run(cmd,
                           stderr=subprocess.PIPE,
                           stdout=devnull)
    output_dict[img_path] = det_result.stderr.decode('utf-8')
    percent_progress = 100*len(output_dict)/len(all_inputs)
    cur_t = perf_counter()
    print(
        f"[multiprocessing generate result] progress: {percent_progress:.1f} %; time elapsed: {cur_t-start_t:.2f} s", end="\r")


if __name__ == "__main__":
    start_t = perf_counter()
    manager = Manager()
    output_dict = manager.dict()
    jobs = []
    all_inputs = []
    for filename in os.listdir(path_to_image):
        if filename.endswith('.jpg') or filename.endswith('.png'):
            all_inputs.append((path_to_image, filename, output_dict))

    with Pool(processes=cpu_count()-2) as pool:
        pool.starmap(process_single_img, all_inputs)

    path_to_output = os.path.join(curpath, sys.argv[2])
    with open(path_to_output, 'a') as fp:
        for img_path, cctag_output in output_dict.items():
            # append 'Processing image' <img_path> to fp
            fp.write('Processing image \"'+img_path+'\"\n')
            fp.write(cctag_output)
    cur_t = perf_counter()
    print(
        f"[multiprocessing generate result] progress: 100 %; time elapsed: {cur_t-start_t:.2f} s")