Block_Extraction.py

import numpy as np
from pathlib import Path
import matplotlib.pyplot as plt

from Functions.Additional_Functions import Line_Length
from Functions.Block_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
    """
    # Plotting 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 Plot_Blocks_Bbox_KeyPoint(Bbox_KeyPoints_Block_LineNumber, Perspective_graph, Points = False, Move_Points = 5, Point_Size = 20, Point_Color = "black"):
    """
    Block Graph Plot Function
    \nInputs:
        \nBbox_KeyPoints_Block_LineNumber(List): A list of lists where each list has format of [Bounding Box, Lines in Bounding Box]
        \nPoints(bool): Boolean parameter to draw key points or not
        \nMove_Points(int): How far each key point should be placed from its original location
        \nPoint_Size(int): Size of each point
        \nPoint_Color(str): Color of points
    \nOutput:
        \n Plot of graph and its bounding boxes 
    """
    # Plotting Graph
    Plot_Graph_KeyPoint(Perspective_graph, Lines = True, Line_Color = 'black', Points = False, Point_Color = 'black')
    # Plotting Bounding Boxes
    for Bb_KP in Bbox_KeyPoints_Block_LineNumber:
        B_Bbox = Bb_KP[0]
        X_Min = B_Bbox[0][0]
        X_Max = B_Bbox[1][0]
        Y_Min = B_Bbox[0][1]
        Y_Max = B_Bbox[1][1]
        X_Bbox = [X_Min, X_Min, X_Max, X_Max, X_Min]
        Y_Bbox = [Y_Min, Y_Max, Y_Max, Y_Min, Y_Min]
        plt.plot(X_Bbox, Y_Bbox, color = "gold", linewidth = 3, alpha=0.5)

        if Points:
            KeyPoints = Bb_KP[1]
            i = 0
            for KP in KeyPoints:
                Line = [KP[0][0], KP[0][1], KP[0][2], KP[0][3]]
                L = Line_Length(Line)
                Normal_Vec1 = (np.array([KP[0][0], KP[0][1]]) - np.array([KP[0][2], KP[0][3]])) / L
                Normal_Vec2 = (np.array([KP[0][2], KP[0][3]]) - np.array([KP[0][0], KP[0][1]])) / L
                Edited_X1 = KP[0][0] + (Normal_Vec2[0] / Move_Points)
                Edited_Y1 = KP[0][1] + (Normal_Vec2[1] / Move_Points)
                Edited_X2 = KP[0][2] + (Normal_Vec1[0] / Move_Points)
                Edited_Y2 = KP[0][3] + (Normal_Vec1[1] / Move_Points)
                plt.scatter([Edited_X1, Edited_X2], [Edited_Y1, Edited_Y2], color = Point_Color, alpha=1, s = Point_Size)
                i += 1

def Block_Extractor(OBJ_file_path, Perspective_Graph, Cam_Distance, Save = False, Save_Path = None):
    """
    Edge Classifire Function
    \nInputs:
        \nOBJ_file_path(str): Path to the .obj file
        \nPerspective_Graph(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)
        \nSave(bool): Boolean parameter to save graph lines as .json file
        \nSave_Path(str): Path to save .json file (if not defind the .json file will be saved in the .obj file location with same name)
    \nOutput:
        \nClassified_Graph(list): A list of lists where each list has format of [[[x1, y1], [x2,  y2]], Edge Type, Edge Color]
    """

    # Loading .obj File
    Roof = trimesh.load(OBJ_file_path)

    # Calculating Parameters for Edge Classification
    Rotation_Angle = 0
    #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)

    Bbox_KeyPoints_Block_LineNumber, _ = Block_Extraction(Roof,
                                                        Cam_Roof,
                                                        Perspective_Graph.tolist(),
                                                        Rotation_Angle)
    
    if Save:
        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]}_Blocks_Graph.json"
            Classified_Edges_Json(Bbox_KeyPoints_Block_LineNumber, File_Name)
        else:
            File_Name = f"{Save_Path}/{OBJ_file_path[len(Save_Path)+1:][:-4]}_Blocks_Graph.json"
            Classified_Edges_Json(Bbox_KeyPoints_Block_LineNumber, File_Name)

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