Update app.py

app.py

@@ -128,71 +128,154 @@ def get_tree(imgpath):
 ROOT = None
 
 
-def display_tree():
-    nodes = []
-    edges = []
-    positions = {}
+# def display_tree():
+#     nodes = []
+#     edges = []
+#     positions = {}
 
-    root = ROOT
+#     root = ROOT
     
-    def traverse(node, depth=0, index=0):
+#     def traverse(node, depth=0, index=0):
 
-        if depth >= 3:
-            return 
+#         if depth >= 3:
+#             return 
             
-        if node not in nodes:
-            nodes.append(node)
-        idx = nodes.index(node)
-        positions[idx] = (depth * 1, index * 1 - len(nodes) / 2)  # Adjusted the multipliers for depth and index
-
-        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, index + len(node.children) / 2)  # Recursively traverse to set positions
-
-    traverse(root)
-    
-    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])
+#         if node not in nodes:
+#             nodes.append(node)
+#         idx = nodes.index(node)
+#         positions[idx] = (depth * 1, index * 1 - len(nodes) / 2)  # Adjusted the multipliers for depth and index
+
+#         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, index + len(node.children) / 2)  # Recursively traverse to set positions
+
+#     traverse(root)
     
-    edge_trace = go.Scatter(
-        x=edge_x, y=edge_y,
-        line=dict(width=2, color='Black'),
-        hoverinfo='none',
-        mode='lines')
+#     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])
     
-    node_x = [pos[0] for pos in positions.values()]
-    node_y = [pos[1] for pos in positions.values()]
+#     edge_trace = go.Scatter(
+#         x=edge_x, y=edge_y,
+#         line=dict(width=2, color='Black'),
+#         hoverinfo='none',
+#         mode='lines')
     
-    node_trace = go.Scatter(
-        x=node_x, y=node_y,
-        mode='markers+text',
-        hoverinfo='text',
-        marker=dict(showscale=False, size=10, color='Goldenrod'),
-        text=[node.name for node in nodes],
-        textposition="top center"
-    )
+#     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=[node.name for node in nodes],
+#         textposition="top center"
+#     )
     
-    layout = go.Layout(
-        title="Tree Visualization",
-        showlegend=False,
-        hovermode='closest',
-        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)
+#     layout = go.Layout(
+#         title="Tree Visualization",
+#         showlegend=False,
+#         hovermode='closest',
+#         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)
+#     )
+    
+#     fig = go.Figure(data=[edge_trace, node_trace], layout=layout)
+#     return fig
+
+
+
+
+
+def create_edge_list(node, edge_list=None, name_map=None):
+    if edge_list is None:
+        edge_list = []
+    if name_map is None:
+        name_map = {}
+    
+    if node.name not in name_map:
+        name_map[node.name] = len(name_map)
+    
+    for child in node.children:
+        if child.name not in name_map:
+            name_map[child.name] = len(name_map)
+        edge_list.append((name_map[node.name], name_map[child.name]))
+        create_edge_list(child, edge_list, name_map)
+    
+    return edge_list, list(name_map.keys())
+
+def display_tree():
+    root = ROOT
+    edge_list, node_names = create_edge_list(root)
+    
+    # Create an igraph Graph from edge list
+    g = ig.Graph(edges=edge_list, directed=True)
+    
+    # Use the Reingold-Tilford tree layout
+    layout = g.layout('rt', root=[0])
+    
+    # Scale the layout to make the tree more compact
+    scale_factor = 0.2  # Adjust this factor as needed
+    layout_coords = [(coord[0] * scale_factor, coord[1] * scale_factor) for coord in layout.coords]
+    
+    x_coords = [coord[0] for coord in layout_coords]
+    y_coords = [-coord[1] for coord in layout_coords]  # invert y-axis for a top-down tree view
+
+    # Create Plotly traces for nodes and edges
+    edge_trace = go.Scatter(
+        x=[None],
+        y=[None],
+        line=dict(width=2, color='#888'),
+        hoverinfo='none',
+        mode='lines'
     )
     
-    fig = go.Figure(data=[edge_trace, node_trace], layout=layout)
-    return fig
+    for v1, v2 in g.get_edgelist():
+        x0, y0 = layout_coords[v1]
+        x1, y1 = layout_coords[v2]
+        # edge_trace['x'] += [x0, x1, None]
+        # edge_trace['y'] += [-y0, -y1, None]  # invert y-axis
 
+        edge_trace['x'] += tuple(list(edge_trace['x']) + [x0, x1, None])
+        edge_trace['y'] += tuple(list(edge_trace['y']) + [-y0, -y1, None])  # invert y-axis
 
+    node_trace = go.Scatter(
+            x=x_coords, y=y_coords,
+            text=node_names,
+            mode='markers+text',
+            hoverinfo='text',
+            textposition='top center',
+            marker=dict(
+                showscale=False,
+                color='Blue',
+                size=10,
+                line_width=2),
+            textfont=dict(
+                size=12,  # Increase the font size as needed
+                color='Black'
+            )
+        )
+
+    # Create a Plotly figure
+    fig = go.Figure(data=[edge_trace, node_trace],
+                    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),
+                        title="Tree Visualization"
+                    ))
+    return fig
 
 
 def get_protoIDs(nodename):