utils/edge_features.py

import numpy as np
import math

import shapely
from shapely.geometry import LineString

import torch
from torch_geometric.utils import to_undirected, remove_self_loops
from torch_geometric.data import Data

def segments(polyline):
    return list(map(LineString, zip(polyline.coords[:-1], polyline.coords[1:])))


def edge_graph(apa_line, apa_wall):
    wall_seg = segments(apa_line)
    num_seg = len(wall_seg)


    edge_lst = []
    for l in range(num_seg):
        seg = wall_seg[l]
        seg_length = seg.length
        seg_pro = apa_wall[l]
        south_cos = wall_segment_cosine("south", seg)
        east_cos = wall_segment_cosine("east", seg)
        north_cos = wall_segment_cosine("north", seg)
        west_cos = wall_segment_cosine("west", seg)

        if south_cos < 0:
            south_cos = 0
        if east_cos < 0:
            east_cos = 0
        if north_cos < 0:
            north_cos = 0
        if west_cos < 0:
            west_cos = 0

        if seg_pro == "I":
            south_cos = 0
            east_cos = 0
            north_cos = 0
            west_cos = 0
            
        if seg_pro == "O":
            seg_boo = 1
        else:
            seg_boo = 0

            

        edge = [seg_boo, seg_length, south_cos, north_cos, west_cos, east_cos]
        edge_lst.append(edge)


    ms_lst = []
    me_lst = []
    for k in range(num_seg):
        if k == (num_seg - 1):
            ms = k
            me = 0
        else:
            ms = k
            me = k+1
        ms_lst.append(ms)
        me_lst.append(me)
    mse = [ms_lst, me_lst]

    datasets = []
    for i in range(2):
        node_features = torch.FloatTensor(edge_lst)
        x = node_features

        edge_index = torch.tensor(mse, dtype=torch.long)
        edge_index, _ = remove_self_loops(edge_index)
        edge_index = to_undirected(edge_index=edge_index)

        data = Data(x=x, edge_index=edge_index)
        datasets.append(data)
    return datasets



def wall_segment_cosine(direction, apa_line_seg):
    seg_s = list(apa_line_seg.coords)[0]
    seg_e = list(apa_line_seg.coords)[1]
    
    normal_x = seg_e[0] - seg_s[0]
    normal_y = seg_e[1] - seg_s[1]
    
    normal_s = (-normal_y, normal_x)
    normal_e = (normal_y, -normal_x)
    
    o = np.array([-normal_y, normal_x])
    w = np.array([normal_y, -normal_x])
    
    if direction == "south":
        d = np.array([-normal_y, normal_x-1])
    if direction == "east":
        d = np.array([-normal_y+1, normal_x])
    if direction == "north":
        d = np.array([-normal_y, normal_x+1])
    if direction == "west":
        d = np.array([-normal_y-1, normal_x])
        
    od = d - o
    ow = w - o
    
    cosine = np.dot(od, ow) / (np.linalg.norm(od) * np.linalg.norm(ow))
    return cosine