def test_finds_singular_outer_path():
mask = np.array([[0, 0, 0, 0, 0], [0, 1, 1, 1, 1], [0, 1, 1, 1, 0], [0, 0, 1, 0, 0]], dtype=np.uint8)
_labels, external_paths, internal_paths = find_contours(mask)
assert len(external_paths) == 1
assert len(internal_paths) == 0
print(mask)
print(np.array(draw_polygon(mask.copy() * 0, external_paths, 1)))
print(external_paths)
assert np.all(mask == draw_polygon(mask.copy() * 0, external_paths, 1))
def test_rectangle_tiny_segments():
square = [
1,
1,
2,
1,
3,
1,
4,
1,
5,
1,
5,
2,
5,
3,
5,
4,
5,
5,
4,
5,
3,
5,
2,
5,
1,
5,
1,
4,
1,
3,
1,
2,
]
expected = np.array(
[
[0, 0, 0, 0, 0, 0, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 1, 1, 1, 1, 1, 0],
[0, 0, 0, 0, 0, 0, 0],
],
dtype=np.uint8,
)
mask = np.zeros((7, 7), dtype=np.int32)
draw_polygon(mask, [square], 1)
assert np.all(mask == expected)
def convert_polygons_to_mask(polygons: List, height: int, width: int, value: Optional[int] = 1) -> np.ndarray:
"""
Converts a list of polygons, encoded as a list of dictionaries into an nd.array mask
Parameters
----------
polygons: list
List of coordinates in the format [{x: x1, y:y1}, ..., {x: xn, y:yn}] or a list of them
as [[{x: x1, y:y1}, ..., {x: xn, y:yn}], ..., [{x: x1, y:y1}, ..., {x: xn, y:yn}]].
Returns
-------
mask: ndarray[float]
ndarray mask of the polygon(s)
"""
sequence = convert_polygons_to_sequences(polygons, height=height, width=width)
mask = np.zeros((height, width)).astype(np.uint8)
draw_polygon(mask, sequence, value)
return mask
def convert_segmentation_to_mask(segmentations: List[List[float]], height: int, width: int):
"""
Converts a polygon represented as a sequence of coordinates into a mask.
Input:
segmentations: list of float values -> [x1, y1, x2, y2, ..., xn, yn]
height: image's height
width: image's width
Output:
torch.tensor
"""
masks = []
for contour in segmentations:
mask = torch.zeros((height, width)).numpy().astype(np.uint8)
masks.append(torch.from_numpy(np.asarray(draw_polygon(mask, contour, 1))))
return torch.stack(masks)
def export(annotation_files: Generator[dt.AnnotationFile, None, None],
output_dir: Path,
mode: str = "grey"):
masks_dir = output_dir / "masks"
masks_dir.mkdir(exist_ok=True, parents=True)
annotation_files = list(annotation_files)
categories = extract_categories(annotation_files)
N = len(categories)
if mode == "index":
if N > 254:
raise ValueError("maximum number of classes supported: 254")
palette = {c: i for i, c in enumerate(categories)}
elif mode == "grey":
if N > 254:
raise ValueError("maximum number of classes supported: 254")
palette = {c: int(i * 255 / (N - 1)) for i, c in enumerate(categories)}
elif mode == "rgb":
if N > 360:
raise ValueError("maximum number of classes supported: 360")
palette = {c: i for i, c in enumerate(categories)}
HSV_colors = [(x / N, 0.8, 1.0) for x in range(N - 1)
] # Generate HSV colors for all classes except for BG
RGB_colors = list(
map(lambda x: [int(e * 255) for e in colorsys.hsv_to_rgb(*x)],
HSV_colors))
RGB_colors.insert(
0, [0, 0, 0]) # Now we add BG class with [0 0 0] RGB value
palette_rgb = {c: rgb for c, rgb in zip(categories, RGB_colors)}
RGB_colors = [c for e in RGB_colors for c in e]
for annotation_file in get_progress_bar(list(annotation_files),
"Processing annotations"):
image_id = os.path.splitext(annotation_file.filename)[0]
outfile = masks_dir / f"{image_id}.png"
outfile.parent.mkdir(parents=True, exist_ok=True)
height = annotation_file.image_height
width = annotation_file.image_width
mask = np.zeros((height, width)).astype(np.uint8)
annotations = [
a for a in annotation_file.annotations
if ispolygon(a.annotation_class)
]
for a in annotations:
cat = a.annotation_class.name
if a.annotation_class.annotation_type == "polygon":
polygon = a.data["path"]
elif a.annotation_class.annotation_type == "complex_polygon":
polygon = a.data["paths"]
sequence = convert_polygons_to_sequences(polygon,
height=height,
width=width)
draw_polygon(mask, sequence, palette[cat])
if mode == "rgb":
mask = Image.fromarray(mask, "P")
mask.putpalette(RGB_colors)
else:
mask = Image.fromarray(mask)
mask.save(outfile)
with open(output_dir / "class_mapping.csv", "w") as f:
f.write(f"class_name,class_color\n")
for c in categories:
if mode == "rgb":
f.write(
f"{c},{palette_rgb[c][0]} {palette_rgb[c][1]} {palette_rgb[c][2]}\n"
)
else:
f.write(f"{c},{palette[c]}\n")
def build_annotation(annotation_file, annotation_id, annotation: dt.Annotation,
categories):
annotation_type = annotation.annotation_class.annotation_type
if annotation_type == "polygon":
sequences = convert_polygons_to_sequences(annotation.data["path"],
rounding=False)
x_coords = [s[0::2] for s in sequences]
y_coords = [s[1::2] for s in sequences]
min_x = np.min([np.min(x_coord) for x_coord in x_coords])
min_y = np.min([np.min(y_coord) for y_coord in y_coords])
max_x = np.max([np.max(x_coord) for x_coord in x_coords])
max_y = np.max([np.max(y_coord) for y_coord in y_coords])
w = max_x - min_x + 1
h = max_y - min_y + 1
# Compute the area of the polygon
poly_area = np.sum([
polygon_area(x_coord, y_coord)
for x_coord, y_coord in zip(x_coords, y_coords)
])
return {
"id": annotation_id,
"image_id": annotation_file.seq,
"category_id": categories[annotation.annotation_class.name],
"segmentation": sequences,
"area": poly_area,
"bbox": [min_x, min_y, w, h],
"iscrowd": 0,
"extra": build_extra(annotation),
}
elif annotation_type == "complex_polygon":
mask = np.zeros(
(annotation_file.image_height, annotation_file.image_width))
sequences = convert_polygons_to_sequences(annotation.data["paths"])
draw_polygon(mask, sequences, 1)
counts = rle_encode(mask)
x_coords = [s[0::2] for s in sequences]
y_coords = [s[1::2] for s in sequences]
min_x = np.min([np.min(x_coord) for x_coord in x_coords])
min_y = np.min([np.min(y_coord) for y_coord in y_coords])
max_x = np.max([np.max(x_coord) for x_coord in x_coords])
max_y = np.max([np.max(y_coord) for y_coord in y_coords])
w = max_x - min_x + 1
h = max_y - min_y + 1
return {
"id": annotation_id,
"image_id": annotation_file.seq,
"category_id": categories[annotation.annotation_class.name],
"segmentation": {
"counts": counts,
"size":
[annotation_file.image_height, annotation_file.image_width]
},
"area": 0,
"bbox": [min_x, min_y, w, h],
"iscrowd": 1,
"extra": build_extra(annotation),
}
elif annotation_type == "tag":
pass
elif annotation_type == "bounding_box":
x = annotation.data["x"]
y = annotation.data["y"]
w = annotation.data["w"]
h = annotation.data["h"]
return build_annotation(
annotation_file,
annotation_id,
dt.make_polygon(
annotation.annotation_class.name,
[{
"x": x,
"y": y
}, {
"x": x + w,
"y": y
}, {
"x": x + w,
"y": y + h
}, {
"x": x,
"y": y + h
}],
),
categories,
)
else:
print(f"skipping unsupported annotation_type '{annotation_type}'")
def test_finds_two_outer_path():
mask = np.array([[0, 0, 0, 0, 0, 0], [0, 1, 1, 0, 1, 0], [0, 1, 1, 0, 1, 0], [0, 0, 1, 0, 1, 0]], dtype=np.uint8,)
_labels, external_paths, internal_paths = find_contours(mask)
assert len(external_paths) == 2
assert len(internal_paths) == 0
assert np.all(mask == draw_polygon(mask.copy() * 0, external_paths, 1))
def export(annotation_files: Iterable[dt.AnnotationFile], output_dir: Path,
mode: str) -> None:
masks_dir: Path = output_dir / "masks"
masks_dir.mkdir(exist_ok=True, parents=True)
annotation_files = list(annotation_files)
categories: List[str] = extract_categories(annotation_files)
num_categories = len(categories)
palette = get_palette(mode=mode, categories=categories)
if mode == "rgb":
# Generate HSV colors for all classes except for BG
HSV_colors = [(x / num_categories, 0.8, 1.0)
for x in range(num_categories - 1)]
RGB_color_list = list(
map(lambda x: [int(e * 255) for e in colorsys.hsv_to_rgb(*x)],
HSV_colors))
# Now we add BG class with [0 0 0] RGB value
RGB_color_list.insert(0, [0, 0, 0])
palette_rgb = {c: rgb for c, rgb in zip(categories, RGB_color_list)}
RGB_colors = [c for e in RGB_color_list for c in e]
for annotation_file in annotation_files:
image_id = os.path.splitext(annotation_file.filename)[0]
outfile = masks_dir / f"{image_id}.png"
outfile.parent.mkdir(parents=True, exist_ok=True)
height = annotation_file.image_height
width = annotation_file.image_width
if height is None or width is None:
raise ValueError(
f"Annotation file {annotation_file.filename} references an image with no height or width"
)
mask: Image.Image = np.zeros((height, width)).astype(np.uint8)
annotations = [
a for a in annotation_file.annotations
if ispolygon(a.annotation_class)
]
for a in annotations:
if isinstance(a, dt.VideoAnnotation):
print(
f"Skipping video annotation from file {annotation_file.filename}"
)
continue
cat = a.annotation_class.name
if a.annotation_class.annotation_type == "polygon":
polygon = a.data["path"]
elif a.annotation_class.annotation_type == "complex_polygon":
polygon = a.data["paths"]
sequence = convert_polygons_to_sequences(polygon,
height=height,
width=width)
draw_polygon(mask, sequence, palette[cat])
if mode == "rgb":
mask = Image.fromarray(mask, "P")
mask.putpalette(RGB_colors)
else:
mask = Image.fromarray(mask)
mask.save(outfile)
with open(output_dir / "class_mapping.csv", "w") as f:
f.write(f"class_name,class_color\n")
for c in categories:
if mode == "rgb":
f.write(
f"{c},{palette_rgb[c][0]} {palette_rgb[c][1]} {palette_rgb[c][2]}\n"
)
else:
f.write(f"{c},{palette[c]}\n")
def test_crop_out_of_bound_vertical_line():
mask = np.zeros((100, 100), dtype=np.int32)
draw_polygon(mask, [[0, -50, 0, -200]], 1)
def test_crops_negative_coordinates():
mask = np.zeros((100, 100), dtype=np.int32)
draw_polygon(mask, [[-50, 0, 50, 50, 200, 200]], 1)
def test_crop_out_of_bound_horizontal_line():
mask = np.zeros((100, 100), dtype=np.int32)
draw_polygon(mask, [[-50, 0, 200, 0]], 1)
def test_writes_the_given_value():
mask_1 = np.zeros((100, 100), dtype=np.int32)
mask_2 = np.zeros((100, 100), dtype=np.int32)
draw_polygon(mask_1, [triangle], 1)
draw_polygon(mask_2, [triangle], 2)
assert np.sum(mask_1) * 2 == np.sum(mask_2)
def test_does_nothing_for_empty_an_empty_polygon():
mask = np.zeros((10, 10), dtype=np.int32)
draw_polygon(mask, [], 1)
assert np.all(mask == 0)
def test_supports_float():
mask = np.zeros(triangle_mask_size, dtype=np.float)
draw_polygon(mask, [triangle], 1)
assert np.all(mask == triangle_result)
def test_out_of_bound():
square = [0, 0, 0, 10, 10, 10, 10, 0]
mask = np.zeros((1, 1), dtype=np.int32)
draw_polygon(mask, [square], 1)
assert np.sum(mask) == 1
def test_decimals_in_path():
square = [0.5, 0.5, 0.5, 10.5, 10.5, 10.5, 10.5, 0.5]
mask = np.zeros((100, 100), dtype=np.int32)
draw_polygon(mask, [square], 1)
assert np.sum(mask) == 11 * 11
