From 423de1be639c855e6589347374b41262c72168e2 Mon Sep 17 00:00:00 2001 From: Andrii Date: Fri, 15 Mar 2024 19:39:31 +0200 Subject: [PATCH] [DAR-1104][External] Fix colours missmatch between files on semanic mask export (#789) --- darwin/datatypes.py | 4 - darwin/exporter/formats/mask.py | 117 +++++++----------- .../exporter/formats/export_mask_test.py | 61 +-------- 3 files changed, 52 insertions(+), 130 deletions(-) diff --git a/darwin/datatypes.py b/darwin/datatypes.py index 386f373c1..10c2e5fc7 100644 --- a/darwin/datatypes.py +++ b/darwin/datatypes.py @@ -1400,7 +1400,6 @@ class MaskTypes: CategoryList = List[str] ExceptionList = List[Exception] UndecodedRLE = List[int] - DecodedRLE = List[List[int]] ColoursDict = Dict[str, int] RgbColors = List[int] HsvColors = List[Tuple[float, float, float]] @@ -1428,7 +1427,6 @@ def validate(self) -> None: @dataclass class RasterLayer: rle: MaskTypes.UndecodedRLE - decoded: MaskTypes.DecodedRLE mask_annotation_ids_mapping: Dict[str, int] slot_names: List[str] = field(default_factory=list) total_pixels: int = 0 @@ -1436,8 +1434,6 @@ class RasterLayer: def validate(self) -> None: if not self.rle: raise ValueError("RasterLayer rle cannot be empty") - if not self.decoded: - raise ValueError("RasterLayer decoded cannot be empty") if not self.mask_annotation_ids_mapping: raise ValueError("RasterLayer mask_annotation_ids_mapping cannot be empty") if not self.slot_names: diff --git a/darwin/exporter/formats/mask.py b/darwin/exporter/formats/mask.py index 0ae38e2c2..cd4672a2e 100644 --- a/darwin/exporter/formats/mask.py +++ b/darwin/exporter/formats/mask.py @@ -162,44 +162,8 @@ def get_render_mode(annotations: List[dt.AnnotationLike]) -> dt.MaskTypes.TypeOf ) -def colours_in_rle( - colours: dt.MaskTypes.ColoursDict, - raster_layer: dt.RasterLayer, - mask_lookup: Dict[str, dt.AnnotationMask], -) -> dt.MaskTypes.ColoursDict: - """ - Returns a dictionary of colours for each mask in the given RLE. - - Parameters - ---------- - colours: dt.MaskTypes.ColoursDict - A dictionary of colours for each mask in the given RLE. - mask_annotations: List[dt.AnnotationMask] - A list of masks to get the colours for. - mask_lookup: Set[str, dt.AnnotationMask] - A lookup table for the masks. - - Returns - ------- - dt.MaskTypes.ColoursDict - A dictionary of colours for each mask in the given RLE. - """ - for uuid, colour_value in raster_layer.mask_annotation_ids_mapping.items(): - mask: Optional[dt.AnnotationMask] = mask_lookup.get(uuid) - - if mask is None: - raise ValueError( - f"Could not find mask with uuid {uuid} in mask lookup table." - ) - - if mask.name not in colours: - colours[mask.name] = colour_value - - return colours # Returns same item as the outset, technically not needed, but best practice. - - -def rle_decode(rle: dt.MaskTypes.UndecodedRLE) -> List[int]: - """Decodes a run-length encoded list of integers. +def rle_decode(rle: dt.MaskTypes.UndecodedRLE, label_colours: Dict[int, int]) -> List[int]: + """Decodes a run-length encoded list of integers and substitutes labels by colours. Args: rle (List[int]): A run-length encoded list of integers. @@ -212,7 +176,7 @@ def rle_decode(rle: dt.MaskTypes.UndecodedRLE) -> List[int]: output = [] for i in range(0, len(rle), 2): - output += [rle[i]] * rle[i + 1] + output += [label_colours[rle[i]]] * rle[i + 1] return output @@ -334,7 +298,7 @@ def render_raster( colours: dt.MaskTypes.ColoursDict, categories: dt.MaskTypes.CategoryList, annotations: List[dt.AnnotationLike], - annotation_file: dt.AnnotationFile, # Not used, but kept for consistency + annotation_file: dt.AnnotationFile, height: int, width: int, ) -> dt.MaskTypes.RendererReturn: @@ -363,16 +327,15 @@ def render_raster( errors: List[Exception] = [] mask_annotations: List[dt.AnnotationMask] = [] - raster_layer: Optional[dt.RasterLayer] = None + raster_layer: dt.RasterLayer - mask_lookup: Dict[str, dt.AnnotationMask] = {} + mask_colours: Dict[str, int] = {} + label_colours: Dict[int, int] = {0: 0} for a in annotations: if isinstance(a, dt.VideoAnnotation): continue - data = a.data - if a.annotation_class.annotation_type == "mask" and a.id: new_mask = dt.AnnotationMask( id=a.id, @@ -387,44 +350,56 @@ def render_raster( mask_annotations.append(new_mask) - if new_mask.id not in mask_lookup: - mask_lookup[new_mask.id] = new_mask - # Add the category to the list of categories if new_mask.name not in categories: categories.append(new_mask.name) - if a.annotation_class.annotation_type == "raster_layer" and (rl := data): - if raster_layer: - errors.append( - ValueError(f"Annotation {a.id} has more than one raster layer") - ) - break + colour_to_draw = categories.index(new_mask.name) + 1 + + if new_mask.id not in mask_colours: + mask_colours[new_mask.id] = colour_to_draw - new_rl = dt.RasterLayer( - rle=rl["dense_rle"], - decoded=rle_decode(rl["dense_rle"]), # type: ignore - slot_names=a.slot_names, - mask_annotation_ids_mapping=rl["mask_annotation_ids_mapping"], - total_pixels=rl["total_pixels"], - ) - new_rl.validate() - raster_layer = new_rl + if new_mask.name not in colours: + colours[new_mask.name] = colour_to_draw + + + raster_layer_list = [a for a in annotations if a.annotation_class.annotation_type == "raster_layer"] - if not raster_layer: - errors.append(ValueError("Annotation has no raster layer")) + if len(raster_layer_list) == 0: + errors.append(ValueError(f"File {annotation_file.filename} has no raster layer")) return errors, mask, categories, colours - if not mask_annotations: - errors.append(ValueError("Annotation has no masks")) + if len(raster_layer_list) > 1: + errors.append( + ValueError(f"File {annotation_file.filename} has more than one raster layer") + ) + return errors, mask, categories, colours + + rl = raster_layer_list[0] + if isinstance(rl, dt.VideoAnnotation): return errors, mask, categories, colours + + raster_layer = dt.RasterLayer( + rle=rl.data["dense_rle"], + slot_names=a.slot_names, + mask_annotation_ids_mapping=rl.data["mask_annotation_ids_mapping"], + total_pixels=rl.data["total_pixels"], + ) + raster_layer.validate() + + for uuid, label in raster_layer.mask_annotation_ids_mapping.items(): + colour_to_draw = mask_colours.get(uuid) - try: - colours = colours_in_rle(colours, raster_layer, mask_lookup) - except Exception as e: - errors.append(e) + if colour_to_draw is None: + errors.append( + ValueError(f"Could not find mask with uuid {uuid} among masks in the file {annotation_file.filename}.") + ) + return errors, mask, categories, colours + + label_colours[label] = colour_to_draw - mask = np.array(raster_layer.decoded, dtype=np.uint8).reshape(height, width) + decoded = rle_decode(raster_layer.rle, label_colours) + mask = np.array(decoded, dtype=np.uint8).reshape(height, width) return errors, mask, categories, colours diff --git a/tests/darwin/exporter/formats/export_mask_test.py b/tests/darwin/exporter/formats/export_mask_test.py index 4e932d2d7..2ddbe004a 100644 --- a/tests/darwin/exporter/formats/export_mask_test.py +++ b/tests/darwin/exporter/formats/export_mask_test.py @@ -18,7 +18,6 @@ from darwin import datatypes as dt from darwin.exporter.formats.mask import ( - colours_in_rle, export, get_or_generate_colour, get_palette, @@ -212,61 +211,17 @@ def test_get_render_mode_raises_value_error_when_no_renderable_annotations_found get_render_mode([dt.Annotation(dt.AnnotationClass("class_3", "invalid"), data={"line": "data"})]) # type: ignore -# Test colours_in_rle -@pytest.fixture -def colours() -> dt.MaskTypes.ColoursDict: - return {"mask1": 1, "mask2": 2} - - -@pytest.fixture -def raster_layer() -> dt.RasterLayer: - return dt.RasterLayer([], [], mask_annotation_ids_mapping={"uuid1": 3, "uuid2": 4}) - - -@pytest.fixture -def mask_lookup() -> Dict[str, dt.AnnotationMask]: - return { - "uuid1": dt.AnnotationMask("mask3", name="mask3"), - "uuid2": dt.AnnotationMask("mask3", name="mask4"), - } - - -def test_colours_in_rle_returns_expected_dict( - colours: dt.MaskTypes.ColoursDict, - raster_layer: dt.RasterLayer, - mask_lookup: Dict[str, dt.AnnotationMask], -) -> None: - expected_dict = {"mask1": 1, "mask2": 2, "mask3": 3, "mask4": 4} - assert colours_in_rle(colours, raster_layer, mask_lookup) == expected_dict - - -def test_colours_in_rle_raises_value_error_when_mask_not_in_lookup( - colours: dt.MaskTypes.ColoursDict, - raster_layer: dt.RasterLayer, - mask_lookup: Dict[str, dt.AnnotationMask], -) -> None: - with pytest.raises(ValueError): - colours_in_rle( - colours, - raster_layer, - { - "uuid9": dt.AnnotationMask("9", name="mask9"), - "uuid10": dt.AnnotationMask("10", name="mask10"), - "uuid11": dt.AnnotationMask("11", name="mask11"), - }, - ) - - # Test RLE decoder def test_rle_decoder() -> None: predication = [1, 2, 3, 4, 5, 6] - expectation = [1, 1, 3, 3, 3, 3, 5, 5, 5, 5, 5, 5] + label_colours = {1: 1, 3: 2, 5: 3} + expectation = [1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3] - assert rle_decode(predication) == expectation + assert rle_decode(predication, label_colours) == expectation odd_number_of_integers = [1, 2, 3, 4, 5, 6, 7] with pytest.raises(ValueError): - rle_decode(odd_number_of_integers) + rle_decode(odd_number_of_integers, label_colours) def test_beyond_polygon_beyond_window() -> None: @@ -503,8 +458,7 @@ def test_render_raster() -> None: dt.AnnotationClass("__raster_layer__", "raster_layer"), { "dense_rle": "my_rle_data", - "decoded": rle_code, - "mask_annotation_ids_mapping": {"mask1": 0, "mask2": 1, "mask3": 2}, + "mask_annotation_ids_mapping": {"mask1": 5, "mask2": 6, "mask3": 7}, "total_pixels": 10000, }, slot_names=["slot1"], @@ -518,11 +472,8 @@ def test_render_raster() -> None: filename="test.txt", ) - with patch("darwin.exporter.formats.mask.rle_decode") as mock_rle_decode, patch( - "darwin.exporter.formats.mask.colours_in_rle" - ) as mock_colours_in_rle: + with patch("darwin.exporter.formats.mask.rle_decode") as mock_rle_decode: mock_rle_decode.return_value = rle_code - mock_colours_in_rle.return_value = {"mask1": 1, "mask2": 2, "mask3": 3} errors, result_mask, result_categories, result_colours = render_raster( mask, colours, categories, annotations, annotation_file, 100, 100