app3.py

import gradio as gr
import plotly.graph_objects as go
import os
from collections import defaultdict
import igraph as ig


# print(os.pwd())

species_to_imgpath = {'bird': './descendent_specific_topk_heatmap_withbb_ep=last_024+051'}

# this has to be there for each species
imgname_to_filepath = {} # this ignores the extension such as .png
nodename_to_protoIDs = defaultdict(list)

for species, imgpath in species_to_imgpath.items():
    for foldername in os.listdir(imgpath):
        if os.path.isdir(os.path.join(imgpath, foldername)):
            folderpath = os.path.join(imgpath, foldername)
            for filename in os.listdir(folderpath):
                if filename.endswith('png') or filename.endswith('jpg'):
                    filepath = os.path.join(folderpath, filename)
                    imgname_to_filepath[filename] = filepath
                    nodename = filename.split('.')[0].split('-')[0]
                    protoID = filename.split('.')[0].split('-')[1]
                    nodename_to_protoIDs[nodename].append(protoID)


class Node():
    def __init__(self, name):
        self.name = name
        self.parent = None
        self.children = [] # list of type Node

    def add_child(child):
            self.children.append(child)


name_to_node = {}

def get_root(node):
    root = node
    while node:
        root = node
        node = node.parent

    return root


def get_tree(imgpath):
    
    for foldername in os.listdir(imgpath):
        if os.path.isdir(os.path.join(imgpath, foldername)):
            folderpath = os.path.join(imgpath, foldername)
            node_name = foldername
            child_names = [filename.split('.')[0].split('-')[0] for filename in os.listdir(folderpath)]

            if node_name in name_to_node:
                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:
                if child_name in name_to_node:
                    child_node = name_to_node[child_name]
                else:
                    child_node = Node(child_name)
                    name_to_node[child_name] = child_node
                child_node.parent = node
                child_nodes.append(child_node)

            node.children = child_nodes
            
        # To be finished
    return get_root(node)
                 

# 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 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

ROOT = None


# def display_tree():
#     nodes = []
#     edges = []
#     positions = {}

#     root = ROOT
    
#     def traverse(node, depth=0, index=0):

#         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])
    
#     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.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)
#     )
    
#     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'
    )
    
    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):
    return gr.Dropdown(choices=nodename_to_protoIDs[nodename], interactive=True)


def get_image(nodename, protoID):
    imgname = '-'.join([nodename, protoID]) + '.png'
    filepath = imgname_to_filepath[imgname]
    return gr.Image(filepath)
    

with gr.Blocks() as demo:

    imgpath = species_to_imgpath['bird']
    print(imgpath)
    ROOT = get_tree(imgpath)
    print(ROOT.name)
    gr.Markdown("## Interactive Tree and Image Display")
    with gr.Row():
        tree_output = gr.Plot(display_tree)  # Connect the function directly
    
    with gr.Row():
        with gr.Column():
            dropdown_1_nodename = gr.Dropdown(label="Select a node name", choices=list(nodename_to_protoIDs.keys()))
            dropdown_1_protos = gr.Dropdown(label="Select a prototype ID", choices=[], allow_custom_value=True)
            image_output_1 = gr.Image()
        with gr.Column():
            dropdown_2_nodename = gr.Dropdown(label="Select a node name", choices=list(nodename_to_protoIDs.keys()))
            dropdown_2_protos = gr.Dropdown(label="Select a prototype ID", choices=[], allow_custom_value=True)
            image_output_2 = gr.Image()

        dropdown_1_nodename.change(get_protoIDs, dropdown_1_nodename, dropdown_1_protos)
        dropdown_1_protos.change(get_image, [dropdown_1_nodename, dropdown_1_protos], image_output_1)
        dropdown_2_nodename.change(get_protoIDs, dropdown_2_nodename, dropdown_2_protos)
        dropdown_2_protos.change(get_image, [dropdown_2_nodename, dropdown_2_protos], image_output_2)
        

# Initialize with placeholder images
# image_output_1.update(display_image_based_on_dropdown_1)
# image_output_2.update(display_image_based_on_dropdown_2)

demo.launch()