Update app.py

app.py

@@ -44,6 +44,7 @@ def get_root(node):
 
     return root
 
+
 def get_tree(imgpath):
     
     for foldername in os.listdir(imgpath):
@@ -56,6 +57,7 @@ def get_tree(imgpath):
                 node = name_to_node[node_name]
             else:
                 node = Node(node_name)
+                name_to_node[node_name] = node
                 
             child_nodes = []
             for child_name in child_names:
@@ -71,22 +73,81 @@ def get_tree(imgpath):
             
         # To be finished
     return get_root(node)
-            
-
-        
-                    
+                 
 
-def display_tree():
-    # This function should create and return a Plotly figure of the tree
+# def display_tree():
+#     # This function should create and return a Plotly figure of the tree
 
-    # Define the nodes and edges for the graph
-    nodes = ['Node 1', 'Node 2', 'Node 3', 'Node 4']
-    edges = [(0, 1), (0, 2), (2, 3)]  # Edges are tuples of node indices
+#     # Define the nodes and edges for the graph
+#     nodes = ['Node 1', 'Node 2', 'Node 3', 'Node 4']
+#     edges = [(0, 1), (0, 2), (2, 3)]  # Edges are tuples of node indices
+    
+#     # Define positions for the nodes (you can use a layout algorithm for more complex graphs)
+#     positions = [(0, 0), (1, 2), (1, -2), (2, 0)]
+    
+#     # Create traces for nodes and edges
+#     edge_x = []
+#     edge_y = []
+#     for edge in edges:
+#         x0, y0 = positions[edge[0]]
+#         x1, y1 = positions[edge[1]]
+#         edge_x.extend([x0, x1, None])
+#         edge_y.extend([y0, y1, None])
+    
+#     edge_trace = go.Scatter(
+#         x=edge_x, y=edge_y,
+#         line=dict(width=2, color='Black'),
+#         hoverinfo='none',
+#         mode='lines')
+    
+#     node_x = [pos[0] for pos in positions]
+#     node_y = [pos[1] for pos in positions]
+    
+#     node_trace = go.Scatter(
+#         x=node_x, y=node_y,
+#         mode='markers+text',
+#         hoverinfo='text',
+#         marker=dict(showscale=False, size=10, color='Goldenrod'),
+#         text=nodes,
+#         textposition="top center"
+#     )
+    
+#     # Define the layout of the graph
+#     layout = go.Layout(
+#         showlegend=False,
+#         hovermode='closest',
+#         margin=dict(b=0, l=0, r=0, t=0),
+#         xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
+#         yaxis=dict(showgrid=False, zeroline=False, showticklabels=False)
+#     )
+    
+#     # Create the figure
+#     fig = go.Figure(data=[edge_trace, node_trace], layout=layout)
+#     return fig
+
+def display_tree(root):
+    nodes = []
+    edges = []
+    positions = {}
     
-    # Define positions for the nodes (you can use a layout algorithm for more complex graphs)
-    positions = [(0, 0), (1, 2), (1, -2), (2, 0)]
+    def traverse(node, depth=0, index=0):
+        # This assigns a unique index to each node, which we can use for positioning
+        if node not in nodes:
+            nodes.append(node)
+        idx = nodes.index(node)
+        positions[idx] = (depth, -len(nodes) + index * 2)  # stagger the y-positions for clarity
+
+        # Loop through each child to connect with edges and calculate their positions
+        for child in node.children:
+            if child not in nodes:
+                nodes.append(child)
+            child_idx = nodes.index(child)
+            edges.append((idx, child_idx))
+            traverse(child, depth + 1, nodes.index(child) - idx)
+
+    # Start the traversal from the root node
+    traverse(root)
     
-    # Create traces for nodes and edges
     edge_x = []
     edge_y = []
     for edge in edges:
@@ -101,28 +162,27 @@ def display_tree():
         hoverinfo='none',
         mode='lines')
     
-    node_x = [pos[0] for pos in positions]
-    node_y = [pos[1] for pos in positions]
+    node_x = [pos[0] for pos in positions.values()]
+    node_y = [pos[1] for pos in positions.values()]
     
     node_trace = go.Scatter(
         x=node_x, y=node_y,
         mode='markers+text',
         hoverinfo='text',
         marker=dict(showscale=False, size=10, color='Goldenrod'),
-        text=nodes,
+        text=[node.name for node in nodes],
         textposition="top center"
     )
     
-    # Define the layout of the graph
     layout = go.Layout(
+        title="Tree Visualization",
         showlegend=False,
         hovermode='closest',
-        margin=dict(b=0, l=0, r=0, t=0),
+        margin=dict(b=0, l=0, r=0, t=40),
         xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
         yaxis=dict(showgrid=False, zeroline=False, showticklabels=False)
     )
     
-    # Create the figure
     fig = go.Figure(data=[edge_trace, node_trace], layout=layout)
     return fig