hololens/stable-diffusion-webui-depthmap-script

1import os
2import cv2
3import glob
4import numpy as np
5import imageio
6from MiDaS.MiDaS_utils import write_depth
7
8BOOST_BASE = 'BoostingMonocularDepth'
9
10BOOST_INPUTS = 'inputs'
11BOOST_OUTPUTS = 'outputs'
12
13def run_boostmonodepth(img_names, src_folder, depth_folder):
14
15    if not isinstance(img_names, list):
16        img_names = [img_names]
17
18    # remove irrelevant files first
19    clean_folder(os.path.join(BOOST_BASE, BOOST_INPUTS))
20    clean_folder(os.path.join(BOOST_BASE, BOOST_OUTPUTS))
21
22    tgt_names = []
23    for img_name in img_names:
24        base_name = os.path.basename(img_name)
25        tgt_name = os.path.join(BOOST_BASE, BOOST_INPUTS, base_name)
26        os.system(f'cp {img_name} {tgt_name}')
27
28        # keep only the file name here.
29        # they save all depth as .png file
30        tgt_names.append(os.path.basename(tgt_name).replace('.jpg', '.png'))
31
32    os.system(f'cd {BOOST_BASE} && python run.py --Final --data_dir {BOOST_INPUTS}/  --output_dir {BOOST_OUTPUTS} --depthNet 0')
33
34    for i, (img_name, tgt_name) in enumerate(zip(img_names, tgt_names)):
35        img = imageio.imread(img_name)
36        H, W = img.shape[:2]
37        scale = 640. / max(H, W)
38
39        # resize and save depth
40        target_height, target_width = int(round(H * scale)), int(round(W * scale))
41        depth = imageio.imread(os.path.join(BOOST_BASE, BOOST_OUTPUTS, tgt_name))
42        depth = np.array(depth).astype(np.float32)
43        depth = resize_depth(depth, target_width, target_height)
44        np.save(os.path.join(depth_folder, tgt_name.replace('.png', '.npy')), depth / 32768. - 1.)
45        write_depth(os.path.join(depth_folder, tgt_name.replace('.png', '')), depth)
46
47def clean_folder(folder, img_exts=['.png', '.jpg', '.npy']):
48
49    for img_ext in img_exts:
50        paths_to_check = os.path.join(folder, f'*{img_ext}')
51        if len(glob.glob(paths_to_check)) == 0:
52            continue
53        print(paths_to_check)
54        os.system(f'rm {paths_to_check}')
55
56def resize_depth(depth, width, height):
57    """Resize numpy (or image read by imageio) depth map
58
59    Args:
60        depth (numpy): depth
61        width (int): image width
62        height (int): image height
63
64    Returns:
65        array: processed depth
66    """
67    depth = cv2.blur(depth, (3, 3))
68    return cv2.resize(depth, (width, height), interpolation=cv2.INTER_AREA)
69