Roof_Graph_Extraction/Graph_Extraction.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
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
import numpy as np
from pathlib import Path
import matplotlib.pyplot as plt

from Functions.Additional_Functions import *
from Functions.pre_proccessing import preproccess
from Functions.Graph_Extraction_Functions import *

def Plot_Graph_KeyPoint(Graph_KeyPoint, Lines = False, Line_Color = 'white', Points = False, Point_Color = 'white'):
    """
    Graph Plot Function
    \nInputs:
        \nGraph_KeyPoint(List): A list of lists or an array where each row has format of [x1, y1, x2, y2]
        \nLines(bool): Boolean parameter to draw lines or not
        \nLine_Color(str): Color of lines
        \nPoints(bool): Boolean parameter to draw key points or not
        \nPoint_Color(str): Color of points
    \nOutput:
        \n Plot of graph
    """
    for GKP in Graph_KeyPoint:
        X = np.array([GKP[0], GKP[2]]).T
        Y = np.array([GKP[1], GKP[3]]).T
        if Lines:
            plt.plot(X, Y, color = Line_Color, linewidth = 1, alpha=1)
        if Points:
            plt.scatter(X, Y, color = Point_Color, s = 20)

def graph_extraction(OBJ_file_path, Cam_Distance, Perspective = False, Save_Graph = True, Save_Path = None):
    """
    Graph Extraction Function:
    \nInputs:
        \nOBJ_file_path(path): Path to the .obj file
        \nPerspective(bool): Boolean parameter which determine if the function should return perspective or orthogonal graph
        \nSave_Graph(bool): Boolean parameter to save graph lines as .csv file
        \nSave_Path(str): Path to save .csv file (if not defind the .csv file will be saved in the .obj file location with same name)
    \nOutput:
        \nGraph_Lines(array): An array with shape M * 4 (each row shows a line whit format [x1, y1, x2, y2] which are start and end points coordinates)
    """
    # Loading .obj File
    Roof = trimesh.load(OBJ_file_path)

    # Creating 3D File From Camera View
    #Cam_Distance = int(Roof.bounds[1][2]) * 25
    #if Cam_Distance > 160:
    #    Cam_Distance = 160
    Cam_Roof = camera_Mesh(Roof, 30, 1, -1, 100, 0, 0, Cam_Distance)

    Rotation_Angle = 0

    #################### Graph Extraction ####################
    if Perspective:
        Merging_Epsilon = 0.001
        Plane_Distance = 0.001
        # Perspective Graph Extraction
        Pre_Perspective_Graph, _ = Graph_Extraction(Roof, Cam_Roof,
                                                    Merging_Epsilon,
                                                    Plane_Distance,
                                                    Rotation_Angle,
                                                    Perspective = Perspective)
        Pre_Perspective_Graph = remove_extra_walls(Pre_Perspective_Graph)
        List_of_Polygons_Points_Normal = Poly_Point_Normal(Pre_Perspective_Graph)
        Perspective_Graph_Key_Points = KeyPoint_Extraction(List_of_Polygons_Points_Normal)
        Perspective_Roof_Graph = preproccess(Perspective_Graph_Key_Points, 0.1)
        if Save_Graph:
            if Save_Path == None:
                Save_Path = str(Path(OBJ_file_path).parent)
                File_Name = f"{Save_Path}/{OBJ_file_path[len(Save_Path)+1:][:-4]}_Perspective_Graph.csv"
                Graph_CSV_Extraction(Perspective_Roof_Graph, File_Name)
            else:
                File_Name = f"{Save_Path}/{OBJ_file_path[len(Save_Path)+1:][:-4]}_Perspective_Graph.csv"
                Graph_CSV_Extraction(Perspective_Roof_Graph, File_Name)

            print(f"Extracted graph file is save to {File_Name}")

        return Perspective_Roof_Graph

    else:
        Merging_Epsilon = 0.02
        Plane_Distance = 0.01
        Pre_Orthogonal_Graph, _ = Graph_Extraction(Roof, Cam_Roof, Merging_Epsilon, Plane_Distance, Rotation_Angle, Perspective= Perspective)
        List_of_Polygons_Points_Normal = Poly_Point_Normal(Pre_Orthogonal_Graph)
        Graph_Key_Points = KeyPoint_Extraction(List_of_Polygons_Points_Normal)
        Orthogonal_Roof_Graph = preproccess(Graph_Key_Points, 0.1)
        if Save_Graph:
            if Save_Path == None:
                Save_Path = str(Path(OBJ_file_path).parent)
                File_Name = f"{Save_Path}/{OBJ_file_path[len(Save_Path)+1:][:-4]}_Orthogonal_Graph.csv"
                Graph_CSV_Extraction(Orthogonal_Roof_Graph, File_Name)
            else:
                File_Name = f"{Save_Path}/{OBJ_file_path[len(Save_Path)+1:][:-4]}_Orthogonal_Graph.csv"
                Graph_CSV_Extraction(Orthogonal_Roof_Graph, File_Name)

            print(f"Extracted graph file is save to {File_Name}")

        return Orthogonal_Roof_Graph