ghostv2/convert.py at main · dimitribarbot/ghostv2 · GitHub

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
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
import os
from typing import cast, Optional

from simple_parsing import ArgumentParser
from tqdm import tqdm

import cv2

from utils.preprocessing.convert_arguments import ConvertArguments


def get_save_path(root_folder: str, source_image: str, output_folder: str, output_extension: str):
    image_name = os.path.splitext(os.path.basename(source_image))[0]
    relative_path = os.path.relpath(os.path.dirname(source_image), root_folder)
    save_folder = os.path.join(output_folder, relative_path)
    save_path = os.path.join(save_folder, f"{image_name}{output_extension}")
    return save_path


def process_one_image(
    source_image: str,
    save_path: str,
):
    image = cv2.imread(source_image)
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    cv2.imwrite(save_path, image)
    return save_path


def process(
    source_image: Optional[str],
    source_folder: Optional[str],
    output_folder: str,
    output_extension: str,
    overwrite: bool,
):
    if source_image is None and source_folder is None:
        raise ValueError("Arguments 'source_image' and 'source_folder' cannot be both empty.")

    if source_folder is None:
        if not os.path.exists(source_image):
            raise ValueError(f"Arguments 'source_image' {source_image} points to a file that does not exist.")

        root_folder = os.path.dirname(source_image)
        save_path = get_save_path(root_folder, source_image, output_folder, output_extension)
        if overwrite or not os.path.exists(save_path):
            print(f"Processing image {source_image}.")
            save_path = process_one_image(
                source_image,
                save_path,
            )
        else:
            print(f"Not processing image {source_image} as target {save_path} already exists and overwrite is false.")
    else:
        if not os.path.exists(source_folder):
            raise ValueError(f"Arguments 'source_folder' {source_folder} points to a folder that does not exist.")

        print(f"Processing images in folder {source_folder}.")
        print("Counting number of files to process.")
        total = sum([len(list(filter(lambda file: overwrite or not os.path.exists(get_save_path(source_folder, os.path.join(root, file), output_folder, output_extension)), files))) \
                     for root, _, files in os.walk(source_folder)])
        print(f"Number of files to process: {total}.")
        with tqdm(total=total) as pbar:
            for root, _, files in os.walk(source_folder):
                for file in files:
                    source_image = os.path.join(root, file)
                    save_path = get_save_path(source_folder, source_image, output_folder, output_extension)
                    if overwrite or not os.path.exists(save_path):
                        save_path = process_one_image(
                            source_image,
                            save_path,
                        )
                        pbar.update()


def main(args: ConvertArguments):
    process(
        args.source_image,
        args.source_folder,
        args.output_folder,
        args.output_extension,
        args.overwrite,
    )


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_arguments(ConvertArguments, dest="arguments")
    args = cast(ConvertArguments, parser.parse_args().arguments)

    main(args)