Skip to content

Commit

Permalink
[DAR-1104][External] Fix colours missmatch between files on semanic m…
Browse files Browse the repository at this point in the history 
…ask export (#789)
  • Loading branch information
@AndriiKlymchuk
AndriiKlymchuk committed Mar 15, 2024
1 parent 37157db commit 423de1b
Show file tree
Hide file tree
Showing 3 changed files with 52 additions and 130 deletions.
4 changes: 0 additions & 4 deletions darwin/datatypes.py
View file
Expand Up @@ -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]]
Expand Down Expand Up @@ -1428,16 +1427,13 @@ 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

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:
Expand Down
117 changes: 46 additions & 71 deletions darwin/exporter/formats/mask.py
View file
Expand Up @@ -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.
Expand All @@ -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

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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,
Expand All @@ -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

Expand Down
61 changes: 6 additions & 55 deletions tests/darwin/exporter/formats/export_mask_test.py
View file
Expand Up @@ -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,
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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"],
Expand All @@ -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
Expand Down

0 comments on commit 423de1b

Please sign in to comment.