fix building rectangle graph

Author: Olexandr Dubchak

core/logic/graph.py

@@ -105,7 +105,7 @@ def add_to_graph(self, node: Node):
     def remove_from_graph(self, node: Node):
         pass
 
-    def get_node_by_id(self, id:int):
+    def get_node_by_id(self, id:int)->Node:
         index = self._nodes_ids.get(id)
         if index is None:
             return None

core/logic/main.py

@@ -3,7 +3,7 @@
 import cv2
 import numpy as np
 
-from core.logic.rectangle_graph import Point, YoloRectangle, DocumentGraph
+from core.logic.rectangle_graph import Point, YoloRectangle, DocumentGraph, RectangleDocumentGraph
 
 bboxes_example = np.load('res_bboxes_example.npy')
 # center_x, center_y, width, height, class_id, probability
@@ -18,8 +18,8 @@
 rec1, rec2 = recs[0], recs[5]
 
 
-#graph = DocumentGraph(recs)
-#graph.build_rectangle_graph()
+graph = RectangleDocumentGraph(recs)
+graph.build_rectangle_graph()
 
 json_example = json.load(open('v3.json', encoding='utf-8'))
 

core/logic/rectangle_graph.py

@@ -3,7 +3,7 @@
 
 import numpy as np
 
-from core.logic.graph import Node
+from core.logic.graph import Node, DocumentGraph
 
 
 class Point:
@@ -146,18 +146,20 @@ def __init__(self, rectangles: List[Rectangle]):
         self.rectangles = rectangles
         self.n = len(rectangles)
         self.distances = np.zeros(shape=(self.n, self.n))
-        self.head = Node(data=None)
-        self.nodes = [Node(data=rec) for rec in rectangles]
         self.calculate_distances()
+        self.graph = DocumentGraph()
+        for i in range(self.n):
+            node = Node(node_id=i, value=f'{self.rectangles[i].top_left_point} ___ {self.rectangles[i].bottom_right_point}')
+            self.graph.add_to_graph(node)
 
     def build_rectangle_graph(self):
         indeces_of_sorted = np.argsort(self.distances, axis=1)
         for rect_index in range(self.n):
             for dist_index in indeces_of_sorted[rect_index, 1:]:
                 position = self.rectangles[rect_index].get_relative_position(self.rectangles[dist_index])
-                node = self.nodes[rect_index]
-                if not node.get_corresponding_attr(position):
-                    node.set_to_corresponding_attr(position, self.nodes[dist_index])
+                node = self.graph.get_node_by_id(rect_index)
+                other_node = self.graph.get_node_by_id(dist_index)
+                node.set_relation(other_node, position)
 
     def calculate_distances(self):
         for i in range(self.n):