image_to_image.utils.eval_extractor

 1"""
 2Helper Module to change the naming of Physgen inference.
 3
 4Just gets called in the evaluation notebook.
 5
 6Its an top-level working script.
 7"""
 8import os
 9import shutil
10import argparse
11
12if __name__ == "__main__":
13    # Adjust these variables:
14    parser = argparse.ArgumentParser()
15    parser.add_argument("--evaluation_path", type=str, required=True)
16    parser.add_argument("--target_path", type=str, required=True)
17    args = parser.parse_args()
18
19    evaluation_path = args.evaluation_path
20    target_path = args.target_path
21
22
23
24
25    def get_building_name(full_name):
26        """
27        buildings_39_real_B.png -> buildings_39.png
28        buildings_39_fake_B.png -> buildings_39.png
29        buildings_39_input_B.png -> buildings_39.png
30        """
31        splitted = full_name.split('_')
32        return f"{splitted[0]}_{splitted[1]}.png"
33
34    os.makedirs(target_path, exist_ok=True)
35
36    target_real_path = f"{target_path}/real"
37    os.makedirs(target_real_path, exist_ok=True)
38
39    target_pred_path = f"{target_path}/pred"
40    os.makedirs(target_pred_path, exist_ok=True)
41
42    target_osm_path = f"{target_path}/osm"
43    os.makedirs(target_osm_path, exist_ok=True)
44
45    for cur_file_name in os.listdir(evaluation_path):
46        cur_file = os.path.join(evaluation_path, cur_file_name)
47        if os.path.isfile(cur_file):
48            if "real_B" in cur_file_name:
49                shutil.copy(cur_file, 
50                            os.path.join(target_real_path, get_building_name(cur_file_name)))
51                print(f"[info] copied real from '{cur_file}' to '{os.path.join(target_real_path, get_building_name(cur_file_name))}'")
52            elif "fake_B" in cur_file_name:
53                shutil.copy(cur_file, 
54                            os.path.join(target_pred_path, get_building_name(cur_file_name)))
55                print(f"[info] copied pred from '{cur_file}' to '{os.path.join(target_pred_path, get_building_name(cur_file_name))}'")
56            if "real_A" in cur_file_name:
57                shutil.copy(cur_file, 
58                            os.path.join(target_osm_path, get_building_name(cur_file_name)))
59                print(f"[info] copied osm from '{cur_file}' to '{os.path.join(target_osm_path, get_building_name(cur_file_name))}'")