Update app.py

app.py

@@ -0,0 +1,320 @@
+import gradio as gr
+from datasets import load_dataset
+from PIL import Image
+import json
+import torch
+from torchvision import transforms
+
+# Load dataset
+dataset = split_dataset['test'] 
+
+# Set up device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+class Sharpen:
+    def __call__(self, img):
+        return img.filter(ImageFilter.SHARPEN)
+
+def preprocess_image(image):
+    # Convert to PIL Image if it's not already
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+
+    # Apply sharpening
+    sharpen = Sharpen()
+    sharpened_image = sharpen(image)
+
+    return sharpened_image
+
+def perform_inference(image):
+    # Preprocess the image
+    inputs = processor(images=image, return_tensors="pt")
+    pixel_values = inputs.pixel_values.to(device)
+
+    # Prepare decoder input ids
+    batch_size = pixel_values.shape[0]
+    decoder_input_ids = torch.full((batch_size, 1), model.config.decoder_start_token_id, device=device)
+
+    # Generate output
+    outputs = model.generate(
+        pixel_values,
+        decoder_input_ids=decoder_input_ids,
+        max_length=max_length, # + 500, #512,  # Adjust as needed
+        early_stopping=True,
+        pad_token_id=processor.tokenizer.pad_token_id,
+        eos_token_id=processor.tokenizer.eos_token_id,
+        use_cache=True,
+        num_beams=1,
+        bad_words_ids=[[processor.tokenizer.unk_token_id]],
+        return_dict_in_generate=True,
+    )
+
+    # Decode the output
+    decoded_output = processor.batch_decode(outputs.sequences)[0]
+    print("Raw model output:", decoded_output)
+
+    return decoded_output
+
+def display_example(index):
+    example = dataset[index]
+    img = example["image"]
+    return img, None, None
+
+def from_json_like_to_xml_like(data):
+    def parse_nodes(nodes):
+        node_elements = []
+        for node in nodes:
+            label = node["label"]
+            node_elements.append(f'<n id="{node["id"]}">{label}</n>')
+        return "<nodes>\n" + "".join(node_elements) + "\n</nodes>"
+
+    def parse_edges(edges):
+        edge_elements = []
+        for edge in edges:
+            edge_elements.append(f'<e src="{edge["source"]}" tgt="{edge["target"]}"/>')
+        return "<edges>\n" + "".join(edge_elements) + "\n</edges>"
+
+    nodes_xml = parse_nodes(data["nodes"])
+    edges_xml = parse_edges(data["edges"])
+    return nodes_xml + "\n" + edges_xml
+
+
+def reshape_json_data_to_fit_visualize_graph(graph_data):
+    nodes = graph_data["nodes"]
+    edges = graph_data["edges"]
+    transformed_nodes = [
+        {"id": nodes["id"][idx], "label": nodes["label"][idx]}
+        for idx in range(len(nodes["id"]))
+    ]
+    transformed_edges = [
+        {"source": edges["source"][idx], "target": edges["target"][idx], "type": "->"}
+        for idx in range(len(edges["source"]))
+    ]
+    return {"nodes": transformed_nodes, "edges": transformed_edges}
+
+def get_ground_truth(index):
+    example = dataset[index]
+    ground_truth = json.dumps(reshape_json_data_to_fit_visualize_graph(example))
+    ground_truth = from_json_like_to_xml_like(json.loads(ground_truth))
+    print(f'Ground truth sequence: {ground_truth}')
+    return ground_truth
+
+def transform_image(img, index, physics_enabled):
+    # Perform inference
+    sequence = perform_inference(img)
+
+    # Transform the sequence to graph data
+    graph_data = transform_sequence(sequence)
+
+    # Generate the graph visualization
+    graph_html = visualize_graph(graph_data, physics_enabled)
+
+    # 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(graph_data, indent=2)
+
+    return graph_html, json_data, sequence
+
+import re
+from typing import Dict, List, Tuple
+
+def transform_sequence(sequence: str) -> Dict[str, List[Dict[str, str]]]:
+    # Extract nodes and edges
+    nodes_match = re.search(r'<nodes>(.*?)</nodes>', sequence, re.DOTALL)
+    edges_match = re.search(r'<edges>(.*?)</edges>', sequence, re.DOTALL)
+
+    if not nodes_match or not edges_match:
+        raise ValueError("Invalid input sequence: nodes or edges not found")
+
+    nodes_text = nodes_match.group(1)
+    edges_text = edges_match.group(1)
+
+    # Parse nodes
+    nodes = []
+    for node_match in re.finditer(r'<n id="\s*(\d+)">(.*?)</n>', nodes_text):
+        node_id, node_label = node_match.groups()
+        nodes.append({
+            "id": node_id.strip(),
+            "label": node_label.strip()
+        })
+
+    # Parse edges
+    edges = []
+    for edge_match in re.finditer(r'<e src="\s*(\d+)" tgt="\s*(\d+)"/>', edges_text):
+        source, target = edge_match.groups()
+        edges.append({
+            "source": source.strip(),
+            "target": target.strip(),
+            "type": "->" 
+        })
+
+    return {
+        "nodes": nodes,
+        "edges": edges
+    }
+
+# function to visualize the extracted graph
+import json
+from pyvis.network import Network
+
+
+def create_graph(nodes, edges, physics_enabled=True):
+    net = Network(
+        notebook=True,
+        height="100vh",
+        width="100vw",
+        bgcolor="#222222",
+        font_color="white",
+        cdn_resources="remote",
+    )
+
+    for node in nodes:
+        net.add_node(
+            node["id"],
+            label=node["label"],
+            title=node["label"],
+            color="blue" if node["label"] == "OOP" else "green",
+        )
+
+    for edge in edges:
+        net.add_edge(edge["source"], edge["target"], title=edge["type"])
+
+    net.force_atlas_2based(
+        gravity=-50,
+        central_gravity=0.01,
+        spring_length=100,
+        spring_strength=0.08,
+        damping=0.4,
+    )
+
+    options = {
+        "nodes": {"physics": physics_enabled},
+        "edges": {"smooth": True},
+        "interaction": {"hover": True, "zoomView": True},
+        "physics": {
+            "enabled": physics_enabled,
+            "stabilization": {"enabled": True, "iterations": 200},
+        },
+    }
+
+    net.set_options(json.dumps(options))
+    return net
+
+
+def visualize_graph(json_data, physics_enabled=True):
+    if isinstance(json_data, str):
+        data = json.loads(json_data)
+    else:
+        data = json_data
+    nodes = data["nodes"]
+    edges = data["edges"]
+    net = create_graph(nodes, edges, physics_enabled)
+    html = net.generate_html()
+    html = html.replace("'", '"')
+    html = html.replace(
+        '<div id="mynetwork"', '<div id="mynetwork" style="height: 100vh; width: 100%;"'
+    )
+    return f"""<iframe style="width: 100%; height: 100vh; border: none; margin: 0; padding: 0;" srcdoc='{html}'></iframe>"""
+
+def update_physics(json_data, physics_enabled):
+    if json_data is None:
+        return None
+
+    data = json.loads(json_data)
+    graph_html = visualize_graph(data, physics_enabled)
+    graph_html = graph_html.replace('height: 100vh;', 'height: 500px;')
+    return graph_html
+
+
+# function to calculate the graph similarity metrics between the prediction and the ground-truth
+def calculate_and_display_metrics(pred_graph, ground_truth_graph):
+    if pred_graph is None or ground_truth_graph is None:
+        return "Please generate a prediction and ensure a ground truth graph is available."
+
+    #removing the start token from the string
+    pred_graph = pred_graph.replace('<s>', "").replace("<newline>", "\n").replace('src=" ', 'src="').replace('tgt=" ', 'tgt="').replace('<n id=" ', '<n id="')
+    print(f'Prediction: {pred_graph}')
+
+    # Assuming the graphs are in the correct format for the calculate_metrics function
+    metrics = model_module.calculate_metrics([pred_graph], [ground_truth_graph])
+
+    # Format the metrics for display
+    overall_metric = metrics[0][0]
+    detailed_metrics = metrics[1][0]
+
+    # output = f"Overall Metric: {overall_metric:.4f}\n\nDetailed Metrics:\n"
+    output = f"Detailed Metrics:\n"
+
+    for key, value in detailed_metrics.items():
+        output += f"{key}: {value:.4f}\n"
+
+    return output
+
+
+def create_interface():
+    with gr.Blocks() as demo:
+        gr.Markdown("# Knowledge Graph Visualizer with Model Inference")
+
+        with gr.Row():
+            index_slider = gr.Slider(
+                minimum=0,
+                maximum=len(dataset) - 1,
+                step=1,
+                label="Example Index"
+            )
+
+        with gr.Row():
+            image_output = gr.Image(type="pil", label="Image", height=500, interactive=False)
+            graph_output = gr.HTML(label="Knowledge Graph")
+
+        with gr.Row():
+            transform_button = gr.Button("Transform")
+            physics_toggle = gr.Checkbox(label="Enable Physics", value=True)
+
+        with gr.Row():
+            json_output = gr.Code(language="json", label="Graph JSON Data")
+            ground_truth_output = gr.Textbox(visible=False)#gr.JSON(label="Ground Truth Graph", visible=False)
+            predicted_raw_sequence = gr.Textbox(visible=False)
+
+        with gr.Row():
+            metrics_button = gr.Button("Calculate Metrics")
+            metrics_output = gr.Textbox(label="Similarity Metrics", lines=10)
+
+        index_slider.change(
+            fn=display_example,
+            inputs=[index_slider],
+            outputs=[image_output, graph_output, json_output],
+        ).then(
+            fn=get_ground_truth,
+            inputs=[index_slider],
+            outputs=[ground_truth_output],
+        )
+
+        transform_button.click(
+            fn=transform_image,
+            inputs=[image_output, index_slider, physics_toggle],
+            outputs=[graph_output, json_output, predicted_raw_sequence],
+        ).then(
+            fn=calculate_and_display_metrics,
+            inputs=[predicted_raw_sequence, ground_truth_output],
+            outputs=[metrics_output]#gr.Textbox(label="Metrics"),
+        )
+        metrics_button.click(
+            fn=calculate_and_display_metrics,
+            inputs=[predicted_raw_sequence, ground_truth_output],
+            outputs=[metrics_output],
+        )
+        physics_toggle.change(
+            fn=update_physics,
+            inputs=[json_output, physics_toggle],
+            outputs=[graph_output],
+        )
+    return demo
+
+# Create and launch the interface
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch(share=True, debug=True)