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