app.py

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

# 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

def display_tree(root):
    nodes = []
    edges = []
    positions = {}
    
    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)
    
    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 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']
    root = get_tree(imgpath)
    gr.Markdown("## Interactive Tree and Image Display")
    with gr.Row():
        tree_output = gr.Plot(display_tree, input=root)  # 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()