Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[DAR-1104][External] Fix colours mismatch between files on semantic mask export #789

Merged
merged 1 commit into from Mar 15, 2024
Merged

[DAR-1104][External] Fix colours mismatch between files on semantic mask export #789

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
4 changes: 0 additions & 4 deletions darwin/datatypes.py
View file
Open in desktop
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
Open in desktop
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
Open in desktop
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