Create app.py

app.py

@@ -0,0 +1,65 @@
+import gradio as gr
+import json
+import numpy as np
+import datasets
+
+bugs_ds = datasets.load_dataset("asgaardlab/SampleDataset", split="validation")
+
+
+def generate_annotations(image_index):
+    image_index = int(image_index)
+    objects_json = bugs_ds[image_index]["Objects JSON"]
+    objects = json.loads(objects_json)
+
+    segmentation_image_rgb = bugs_ds[image_index]["Segmentation Image"]
+    segmentation_image_rgb = np.array(segmentation_image_rgb)
+
+    annotations = []
+    for obj in objects:
+        color = tuple(obj["color"].values())[:-1]
+        mask = np.all(segmentation_image_rgb == np.array(color), axis=-1).astype(
+            np.float32
+        )
+        annotations.append((mask, obj["name"]))
+
+    object_count = bugs_ds[image_index]["Object Count"]
+    victim_name = bugs_ds[image_index]["Victim Name"]
+    bug_type = bugs_ds[image_index]["Tag"]
+
+    return (
+        (bugs_ds[image_index]["Correct Image"], annotations),
+        objects,
+        object_count,
+        victim_name,
+        bug_type,
+    )
+
+
+# Setting up the Gradio interface using blocks API
+with gr.Blocks() as demo:
+    gr.Markdown(
+        "Enter the image index and click **Submit** to view the segmentation annotations."
+    )
+    with gr.Row():
+        inp = gr.Slider(
+            minimum=0, maximum=len(bugs_ds) - 1, step=1, label="Image Index"
+        )
+        btn = gr.Button("Submit")
+    with gr.Row():
+        with gr.Column():
+            object_count = gr.Number(label="Object Count")
+            victim_name = gr.Textbox(label="Victim Name")
+            bug_type = gr.Textbox(label="Bug Type")
+
+        seg_img = gr.AnnotatedImage()
+
+    with gr.Row():
+        json_data = gr.JSON()
+
+    btn.click(
+        fn=generate_annotations,
+        inputs=inp,
+        outputs=[seg_img, json_data, object_count, victim_name, bug_type],
+    )
+
+demo.launch()