-
Notifications
You must be signed in to change notification settings - Fork 2
/
main.py
57 lines (38 loc) · 1.68 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
import os
import cv2
import argparse
from tqdm import tqdm
from annotators import add_annotators_to_arg_parser
def main():
parser = argparse.ArgumentParser(
description="Batch converts images to depth map using MiDaS"
)
parser.add_argument("input_dir", help="input directory")
parser.add_argument("output_dir", help="output directory")
parser.add_argument("annotator", type=str, default="midas", help="Annotator to use")
parser.add_argument("--res", type=int, default=512,
help="Resolution to process at. -1 for original size (be careful for large images!)")
parser.add_argument("--gpu", type=int, default=0, help="GPU id to use")
add_annotators_to_arg_parser(parser)
args = parser.parse_args()
if args.gpu is not None:
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
# make output directory
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir, exist_ok=True)
from annotators import annotate
img_ext = [".jpg", ".jpeg", ".png", ".webp"]
image_paths = [img_path for img_path in os.listdir(args.input_dir) if
os.path.splitext(img_path)[1].lower() in img_ext]
print(f"Found {len(image_paths)} images")
for img_path in tqdm(image_paths):
full_img_path = os.path.join(args.input_dir, img_path)
img = cv2.imread(full_img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
output = annotate(img, args)
output = output.astype('uint8')
output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
cv2.imwrite(os.path.join(args.output_dir, img_path), output)
print("FIN")
if __name__ == "__main__":
main()