Roof_Graph_Extraction/Data_Create_Block_Detector.py at main · ramintkml/Roof_Graph_Extraction · 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
import warnings
warnings.filterwarnings("ignore")

from tqdm import tqdm
import sys
import glob
import os
import torch
os.chdir("/home/training/Takmil/Roof_Graph_Extraction/Functions")
from Functions.Graph_Extraction_Functions import *
from Functions.Edge_Classification_Functions import Classified_Edges_Json
from Functions.Block_Detector_Model_Functions import *


if __name__ == "__main__":
    Root_Dataset_Path = "/home/training/Takmil/Datasets/Block_Detector"
    Dataset_Path_Exist_Train = os.path.isdir(f"{Root_Dataset_Path}/Train1")
    if not Dataset_Path_Exist_Train:
        os.mkdir(f"{Root_Dataset_Path}/Train1")

    Dataset_Path_Exist_Test = os.path.isdir(f"{Root_Dataset_Path}/Test1")
    if not Dataset_Path_Exist_Train:
        os.mkdir(f"{Root_Dataset_Path}/Test1")

    Pers_Graphs = glob.glob("/home/training/Takmil/Datasets/One_Floor_New/*_Perspective.csv")

    Dataset_Path_Train = f"{Root_Dataset_Path}/Train1"
    Dataset_Path_Test = f"{Root_Dataset_Path}/Test1"
    for idx in tqdm(list(range(500))):
        try:
            Pers_File = Pers_Graphs[idx]
            Ortho_File = Pers_File[:-16] + "_Orthogonal.csv"
            Roof_File = Pers_File[:-16] + ".obj"
            All_Data, Images, Unrotated_Classified_Graphs, Unrotated_Graphs, Features, All_Centers = Block_Detector_Dataloder(Pers_File,
                                                                                                                              Ortho_File,
                                                                                                                              Roof_File)
            # Save the tensor
            if 0 < idx <= 100:
                Unrotated_Classified_Graph = Unrotated_Classified_Graphs
                Unrotated_Graph = Unrotated_Graphs
                data = All_Data
                Classified_File_Name = rf"{Dataset_Path_Test}/Data_{str(idx)}_Graph_Classified.json"
                Classified_Edges_Json(Unrotated_Classified_Graph, Classified_File_Name)
                Graph_File_Name = rf"{Dataset_Path_Test}/Data_{str(idx)}_Graph.csv"
                Graph_CSV_Extraction(Unrotated_Graph, Graph_File_Name)
                Pers_Nodes_File_Name = rf"{Dataset_Path_Test}/Data_{str(idx)}_Nodes.csv"
                Graph_CSV_Extraction(All_Centers.tolist(), Pers_Nodes_File_Name)
                torch.save(data, rf"{Dataset_Path_Test}/Data_{str(idx)}.pt")
                # print(f"Number of created data: {idx}")
            elif 100 < idx:
                Unrotated_Classified_Graph = Unrotated_Classified_Graphs
                Unrotated_Graph = Unrotated_Graphs
                data = All_Data
                Classified_File_Name = rf"{Dataset_Path_Train}/Data_{str(idx)}_Graph_Classified.json"
                Classified_Edges_Json(Unrotated_Classified_Graph, Classified_File_Name)
                Graph_File_Name = rf"{Dataset_Path_Train}/Data_{str(idx)}_Graph.csv"
                Graph_CSV_Extraction(Unrotated_Graph, Graph_File_Name)
                Pers_Nodes_File_Name = rf"{Dataset_Path_Train}/Data_{str(idx)}_Nodes.csv"
                Graph_CSV_Extraction(All_Centers.tolist(), Pers_Nodes_File_Name)
                torch.save(data, rf"{Dataset_Path_Train}/Data_{str(idx)}.pt")
                # print(f"Number of created data: {idx}")
        except:
            pass