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 ROOT = None def display_tree(): nodes = [] edges = [] positions = {} root = ROOT 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'] 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()