Update app.py

app.py

@@ -8,21 +8,20 @@ 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 = bugs_ds[image_index]["Objects JSON (Correct)"]
     objects = json.loads(objects_json)
 
-    segmentation_image_rgb = bugs_ds[image_index]["Segmentation Image"]
+    segmentation_image_rgb = bugs_ds[image_index]["Segmentation Image (Correct)"]
     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"]))
+        color = tuple(obj["color"])[:-1]
+        mask = np.all(segmentation_image_rgb[:, :, :3] == np.array(color), axis=-1).astype(np.float32)
+
+        annotations.append((mask, obj["labelName"]))
 
-    object_count = bugs_ds[image_index]["Object Count"]
+    object_count = 0 # bugs_ds[image_index]["Object Count"]
     victim_name = bugs_ds[image_index]["Victim Name"]
     bug_type = bugs_ds[image_index]["Tag"]