Spaces:
Sleeping
Sleeping
Update app.py
Browse files
app.py
CHANGED
@@ -8,21 +8,20 @@ bugs_ds = datasets.load_dataset("asgaardlab/SampleDataset", split="validation")
|
|
8 |
|
9 |
def generate_annotations(image_index):
|
10 |
image_index = int(image_index)
|
11 |
-
objects_json = bugs_ds[image_index]["Objects JSON"]
|
12 |
objects = json.loads(objects_json)
|
13 |
|
14 |
-
segmentation_image_rgb = bugs_ds[image_index]["Segmentation Image"]
|
15 |
segmentation_image_rgb = np.array(segmentation_image_rgb)
|
16 |
|
17 |
annotations = []
|
18 |
for obj in objects:
|
19 |
-
color = tuple(obj["color"]
|
20 |
-
mask = np.all(segmentation_image_rgb == np.array(color), axis=-1).astype(
|
21 |
-
|
22 |
-
)
|
23 |
-
annotations.append((mask, obj["name"]))
|
24 |
|
25 |
-
object_count = bugs_ds[image_index]["Object Count"]
|
26 |
victim_name = bugs_ds[image_index]["Victim Name"]
|
27 |
bug_type = bugs_ds[image_index]["Tag"]
|
28 |
|
|
|
8 |
|
9 |
def generate_annotations(image_index):
|
10 |
image_index = int(image_index)
|
11 |
+
objects_json = bugs_ds[image_index]["Objects JSON (Correct)"]
|
12 |
objects = json.loads(objects_json)
|
13 |
|
14 |
+
segmentation_image_rgb = bugs_ds[image_index]["Segmentation Image (Correct)"]
|
15 |
segmentation_image_rgb = np.array(segmentation_image_rgb)
|
16 |
|
17 |
annotations = []
|
18 |
for obj in objects:
|
19 |
+
color = tuple(obj["color"])[:-1]
|
20 |
+
mask = np.all(segmentation_image_rgb[:, :, :3] == np.array(color), axis=-1).astype(np.float32)
|
21 |
+
|
22 |
+
annotations.append((mask, obj["labelName"]))
|
|
|
23 |
|
24 |
+
object_count = 0 # bugs_ds[image_index]["Object Count"]
|
25 |
victim_name = bugs_ds[image_index]["Victim Name"]
|
26 |
bug_type = bugs_ds[image_index]["Tag"]
|
27 |
|