Update app.py

app.py

@@ -1,15 +1,21 @@
 import gradio as gr
-from datasets import load_dataset
+from datasets import load_dataset, Dataset
 from PIL import Image
 import io
 import base64
-import json
+import json 
 from graph_visualization import visualize_graph
 
+# branch_name = "edges-sorted-ascending"
+branch_name = "Sorted_edges"
+
 # Load the dataset
-dataset = load_dataset("Zaherrr/OOP_KG_MarkMap_Synthetic_Dataset", split='data')
+# dataset = load_dataset("Zaherrr/OOP_KG_Dataset", split='data', revision=branch_name)
+dataset = load_dataset("Zaherrr/OOP_KG_MarkMap_Synthetic_Dataset", split='data') #, revision=branch_name)
 print(f'This is the dataset: {dataset}')
+
 print(dataset.info)
+
 print(f'This is an example: {dataset[-5]}')
 
 def reshape_json_data_to_fit_visualize_graph(graph_data):
@@ -26,43 +32,33 @@ def reshape_json_data_to_fit_visualize_graph(graph_data):
     graph_data = {"nodes": transformed_nodes, "edges": transformed_edges}
     return graph_data
 
-def image_to_base64(img):
-    """Convert a PIL image to a base64 string."""
-    buffer = io.BytesIO()
-    img.save(buffer, format="PNG")
-    base64_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
-    return base64_str
-
 def display_example(index):
     example = dataset[index]
     img = example["image"]
 
-    # Convert the image to a base64 string
-    base64_image = image_to_base64(img)
-
-    # Create the HTML for zoomable image
-    img_html = f'''
-    <div style="width:100%; height: 500px; overflow: hidden; position:relative;">
-        <img src="data:image/png;base64,{base64_image}" style="width:100%; cursor: zoom-in;" onclick="this.style.cursor = 'zoom-out'; this.style.transform = 'scale(2)';" onmouseleave="this.style.cursor = 'zoom-in'; this.style.transform = 'scale(1)';" />
-    </div>
-    '''
-
+    # Get image dimensions
+    img_width, img_height = img.size
+    
     # Prepare the graph data
     graph_data = {"nodes": example["nodes"], "edges": example["edges"]}
     transformed_graph_data = reshape_json_data_to_fit_visualize_graph(graph_data)
     
     # Generate the graph visualization
     graph_html = visualize_graph(transformed_graph_data)
-    
+
+    # Modify the iframe to have a fixed height
+    graph_html = graph_html.replace('height: 100vh;', 'height: 500px;')
+
     # Convert graph_data to a formatted JSON string
     json_data = json.dumps(transformed_graph_data, indent=2)
     
-    return img_html, graph_html, json_data, transformed_graph_data, f"Width: {img.width}px, Height: {img.height}px"
+    return img, graph_html, json_data, transformed_graph_data, f"Width: {img_width}px, Height: {img_height}px"
 
 def create_interface():
     with gr.Blocks() as demo:
+    
         gr.Markdown("# Knowledge Graph Visualizer for the [Zaherrr/OOP_KG_MarkMap_Synthetic_Dataset](https://huggingface.co/datasets/Zaherrr/OOP_KG_MarkMap_Synthetic_Dataset) dataset")
-
+        
         with gr.Row():
             index_slider = gr.Slider(
                 minimum=0,
@@ -70,9 +66,9 @@ def create_interface():
                 step=1,
                 label="Example Index"
             )
-
+        
         with gr.Row():
-            image_output = gr.HTML(label="Zoomable Image")
+            image_output = gr.Image(type="pil", label="Image", height=500)
             graph_output = gr.HTML(label="Knowledge Graph")
 
         with gr.Row():
@@ -81,7 +77,7 @@ def create_interface():
                 placeholder="Width and Height will appear here",
                 interactive=False,
             )
-
+            
         with gr.Row():
             json_output = gr.Code(language="json", label="Graph JSON Data")
             text_output = gr.Textbox(
@@ -89,16 +85,16 @@ def create_interface():
                 placeholder="Text data will appear here",
                 interactive=False,
             )
-
+        
         index_slider.change(
             fn=display_example,
             inputs=[index_slider],
             outputs=[image_output, graph_output, json_output, text_output, dimensions_output],
         )
-
+    
     return demo
 
 # Create and launch the interface
 if __name__ == "__main__":
     demo = create_interface()
-    demo.launch()
+    demo.launch()