Env-Rec-OSM-WAIRD/run.py at main · YerevaNN/Env-Rec-OSM-WAIRD · 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
import logging
import os
import random
import sys
import warnings

import hydra
import torch
from dotenv import load_dotenv
from omegaconf import DictConfig
from pytorch_lightning import seed_everything

log = logging.getLogger(__name__)

_PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))
if _PROJECT_ROOT not in sys.path:
    sys.path.insert(0, _PROJECT_ROOT)

load_dotenv(os.path.join(_PROJECT_ROOT, ".env"))


@hydra.main(config_path="configs", config_name="train", version_base="1.2")
def main(config: DictConfig) -> None:
    from src import utils, train, evaluate, pred, reconstruction_visualize

    warnings.filterwarnings("ignore", ".*beta state*")

    terminal_col = config.get("terminal_col")
    if terminal_col:
        terminal_row = config.get("terminal_row", 24)
        utils.set_winsize(sys.stdin, terminal_col, terminal_row)
        utils.set_winsize(sys.stderr, terminal_col, terminal_row)
        utils.set_winsize(sys.stdout, terminal_col, terminal_row)

    if config.seed == -1:
        config.seed = random.randint(0, 10 ** 8)

    seed_everything(config.seed)
    log.info(f"Run dir: {os.path.realpath('./')}")

    # Ensure that all operations are deterministic on GPU (if used) for reproducibility
    torch.backends.cudnn.determinstic = True
    torch.backends.cudnn.benchmark = False

    if config.get("print_config"):
        utils.print_config(config, fields=tuple(config.keys()), resolve=True)

    if config.get("ignore_warnings"):
        log.info("Disabling python warnings! <config.ignore_warnings=True>")
        warnings.filterwarnings("ignore")

    if config.name == "train":
        return train(config)

    if config.name == "reconstruction_visualize":
        return reconstruction_visualize(config)

    if config.name == "inference":
        return pred(config)

    if config.name == "evaluation":
        return evaluate(config)

    if config.name == "oslo_generate_data":
        from src.oslo import generate_data
        return generate_data(config)

    if config.name == "oslo_retrofit_info":
        from src.oslo import retrofit_info
        return retrofit_info(config)

    if config.name == "oslo_create_info_json":
        from src.oslo import create_info_json
        return create_info_json(config)

    if config.name == "oslo_train_test_split":
        from src.oslo import train_test_val
        return train_test_val(config)

    if config.name == "oslo_build_full_area":
        from src.oslo import build_full_area
        return build_full_area(config)

    if config.name == "oslo_generate_crop_cache":
        from src.oslo import generate_crop_cache
        return generate_crop_cache(config)

    if config.name == "oslo_stitch_inference":
        from src.oslo import stitch_inference
        return stitch_inference(config)

    if config.name == "oslo_stitch_evaluate":
        from src.oslo import stitch_evaluate
        return stitch_evaluate(config)

    if config.name == "oslo_stitch_visualize":
        from src.oslo_stitch_visualize import oslo_stitch_visualize
        return oslo_stitch_visualize(config)


if __name__ == "__main__":
    hydra.core.global_hydra.GlobalHydra.instance().clear()
    try:
        main()
    finally:
        hydra.core.global_hydra.GlobalHydra.instance().clear()