ciagen/run.py at main · fennecinspace/ciagen · 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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
import os

import hydra
from omegaconf import DictConfig


def help_task(cfg: DictConfig) -> None:
    title = "Controllable Image Augmentation. Using StableDiffusion + ControlNet."
    sep = "=" * len(title)

    lines = [
        sep,
        title,
        "",
        "Possible tasks:",
        str([", ".join(list(allowed_tasks.keys())[1:])]),
        "",
        "For basic usage do `python run.py task=<task>`.",
        "You can also modify the `config.yaml` file.",
        "For more information, see the README.md file.",
        sep,
    ]
    maximum_length = max(len(line) for line in lines)
    new_lines = [f"| {line} {' ' * (maximum_length - len(line))}|\n" for line in lines]
    print("".join(new_lines))


def gen_task(cfg: DictConfig) -> None:
    from ciagen.hydra_compat import run_gen

    run_gen(cfg)


def dtd_task(cfg: DictConfig) -> None:
    from ciagen.hydra_compat import run_dtd

    run_dtd(cfg)


def ptd_task(cfg: DictConfig) -> None:
    from ciagen.hydra_compat import run_ptd

    run_ptd(cfg)


def filtering_task(cfg: DictConfig) -> None:
    from ciagen.hydra_compat import run_filtering

    run_filtering(cfg)


def auto_captioner_task(cfg: DictConfig) -> None:
    from ciagen.api.caption import caption
    from ciagen.data.paths import generate_all_paths

    paths = generate_all_paths(cfg)
    caption(
        images=paths["real_images"],
        captions_dir=paths["real_captions"],
        engine=cfg["auto_captioner"]["service"]["engine"],
        model=cfg["auto_captioner"]["service"]["model"],
        api_key=cfg["auto_captioner"]["service"].get("api_key"),
        image_formats=list(cfg["data"]["image_formats"]),
    )


def prepare_data_task(cfg: DictConfig) -> None:
    from ciagen.data.paths import generate_all_paths

    dataset = cfg["data"]["base"]
    paths = generate_all_paths(cfg)

    if dataset == "coco":
        from examples.prepare_coco import prepare_coco

        prepare_coco(cfg, paths)
    elif dataset == "flickr30k":
        from examples.prepare_flickr30k import prepare_flickr30k

        prepare_flickr30k(cfg, paths)
    elif dataset in ("fer_real", "fer_gen_1_5", "fer_gen_2_1"):
        from examples.prepare_fer import prepare_fer

        prepare_fer(cfg, paths)
    elif dataset == "mocs":
        from examples.prepare_mocs import prepare_mocs

        prepare_mocs(cfg, paths)
    else:
        print(f"[ERROR] Unknown dataset: {dataset}")


def mix_dataset(cfg: DictConfig) -> None:
    from ciagen.data.paths import generate_all_paths

    dataset = cfg["data"]["base"]
    paths = generate_all_paths(cfg)
    if dataset in ("coco", "flickr30k"):
        from examples.mix_yolo_dataset import mix_yolo

        mix_yolo(cfg, paths)
    elif dataset == "fer":
        from examples.mix_fer_dataset import mix_fer

        mix_fer(cfg, paths)
    else:
        raise ValueError(f"Unknown dataset: {dataset}")


def train(cfg: DictConfig) -> None:
    from ciagen.data.paths import generate_all_paths

    dataset = cfg["data"]["base"]
    paths = generate_all_paths(cfg)
    if dataset in ("coco", "flickr30k"):
        from examples.train_yolo import train_yolo

        train_yolo(cfg, paths)
    elif dataset == "fer":
        from examples.train_classifier import train_classifier

        train_classifier(cfg, paths)
    else:
        raise ValueError(f"Unknown dataset: {dataset}")


allowed_tasks = {
    "help": help_task,
    "prepare_data": prepare_data_task,
    "auto_caption": auto_captioner_task,
    "gen": gen_task,
    "dtd": dtd_task,
    "ptd": ptd_task,
    "filtering": filtering_task,
    "mix": mix_dataset,
    "train": train,
}

path_to_config = os.path.join(os.getcwd(), "ciagen", "conf")


@hydra.main(version_base=None, config_path=path_to_config, config_name="config")
def my_app(cfg: DictConfig) -> None:
    task = cfg["task"]
    if task not in allowed_tasks:
        raise ValueError(f"Unknown task: {task}. Available: {list(allowed_tasks.keys())}")

    allowed_tasks[task](cfg)


if __name__ == "__main__":
    my_app()