Spaces:
No application file
No application file
File size: 3,126 Bytes
430de99 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import json
import numpy as np
import os
import tempfile
import unittest
import pycocotools.mask as mask_util
from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.data.datasets.coco import convert_to_coco_dict, load_coco_json
from detectron2.structures import BoxMode
def make_mask():
"""
Makes a donut shaped binary mask.
"""
H = 100
W = 100
mask = np.zeros([H, W], dtype=np.uint8)
for x in range(W):
for y in range(H):
d = np.linalg.norm(np.array([W, H]) / 2 - np.array([x, y]))
if d > 10 and d < 20:
mask[y, x] = 1
return mask
def uncompressed_rle(mask):
l = mask.flatten(order="F").tolist()
counts = []
p = False
cnt = 0
for i in l:
if i == p:
cnt += 1
else:
counts.append(cnt)
p = i
cnt = 1
counts.append(cnt)
return {"counts": counts, "size": [mask.shape[0], mask.shape[1]]}
def make_dataset_dicts(mask, compressed: bool = True):
"""
Returns a list of dicts that represents a single COCO data point for
object detection. The single instance given by `mask` is represented by
RLE, either compressed or uncompressed.
"""
record = {}
record["file_name"] = "test"
record["image_id"] = 0
record["height"] = mask.shape[0]
record["width"] = mask.shape[1]
y, x = np.nonzero(mask)
if compressed:
segmentation = mask_util.encode(np.asarray(mask, order="F"))
else:
segmentation = uncompressed_rle(mask)
min_x = np.min(x)
max_x = np.max(x)
min_y = np.min(y)
max_y = np.max(y)
obj = {
"bbox": [min_x, min_y, max_x, max_y],
"bbox_mode": BoxMode.XYXY_ABS,
"category_id": 0,
"iscrowd": 0,
"segmentation": segmentation,
}
record["annotations"] = [obj]
return [record]
class TestRLEToJson(unittest.TestCase):
def test(self):
# Make a dummy dataset.
mask = make_mask()
DatasetCatalog.register("test_dataset", lambda: make_dataset_dicts(mask))
MetadataCatalog.get("test_dataset").set(thing_classes=["test_label"])
# Dump to json.
json_dict = convert_to_coco_dict("test_dataset")
with tempfile.TemporaryDirectory() as tmpdir:
json_file_name = os.path.join(tmpdir, "test.json")
with open(json_file_name, "w") as f:
json.dump(json_dict, f)
# Load from json.
dicts = load_coco_json(json_file_name, "")
# Check the loaded mask matches the original.
anno = dicts[0]["annotations"][0]
loaded_mask = mask_util.decode(anno["segmentation"])
self.assertTrue(np.array_equal(loaded_mask, mask))
def test_uncompressed_RLE(self):
mask = make_mask()
rle = mask_util.encode(np.asarray(mask, order="F"))
uncompressed = uncompressed_rle(mask)
compressed = mask_util.frPyObjects(uncompressed, *rle["size"])
self.assertEqual(rle, compressed)
|