Skip to content

Commit 3c1d84a

Browse files
arthur-encordarthur-encord
committed
Add bitmask tests
1 parent 12ba12a commit 3c1d84a

File tree

2 files changed

+363
-4
lines changed

2 files changed

+363
-4
lines changed

‎encord/objects/ontology_labels_impl.py‎

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1384,33 +1384,49 @@ def frame(self) -> int:
13841384

13851385
@property
13861386
def width(self) -> int:
1387-
"""Get the width of the image or image group.
1387+
"""Get the width of the frame.
13881388
13891389
Returns:
13901390
int: The width of the image or image group.
13911391
13921392
Raises:
13931393
LabelRowError: If the width is not set for the data type.
13941394
"""
1395-
if self._label_row.data_type in [DataType.IMG_GROUP]:
1395+
if self._label_row.data_type == DataType.IMG_GROUP:
13961396
return self._frame_level_data().width
1397+
elif self._label_row.data_type == DataType.DICOM:
1398+
frame_metadata = self._label_row._frame_metadata[self._frame]
1399+
if frame_metadata is not None:
1400+
return frame_metadata.width
1401+
elif self._label_row_read_only_data.width is not None:
1402+
return self._label_row_read_only_data.width
1403+
else:
1404+
raise LabelRowError(f"Width is expected but not set for the data type {self._label_row.data_type}")
13971405
elif self._label_row_read_only_data.width is not None:
13981406
return self._label_row_read_only_data.width
13991407
else:
14001408
raise LabelRowError(f"Width is expected but not set for the data type {self._label_row.data_type}")
14011409

14021410
@property
14031411
def height(self) -> int:
1404-
"""Get the height of the image or image group.
1412+
"""Get the height of the frame.
14051413
14061414
Returns:
14071415
int: The height of the image or image group.
14081416
14091417
Raises:
14101418
LabelRowError: If the height is not set for the data type.
14111419
"""
1412-
if self._label_row.data_type in [DataType.IMG_GROUP]:
1420+
if self._label_row.data_type == DataType.IMG_GROUP:
14131421
return self._frame_level_data().height
1422+
elif self._label_row.data_type == DataType.DICOM:
1423+
frame_metadata = self._label_row._frame_metadata[self._frame]
1424+
if frame_metadata is not None:
1425+
return frame_metadata.height
1426+
elif self._label_row_read_only_data.height is not None:
1427+
return self._label_row_read_only_data.height
1428+
else:
1429+
raise LabelRowError(f"Height is expected but not set for the data type {self._label_row.data_type}")
14141430
elif self._label_row_read_only_data.height is not None:
14151431
return self._label_row_read_only_data.height
14161432
else:

‎tests/objects/test_bitmask_validation.py‎

Lines changed: 343 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,343 @@
1+
from datetime import datetime
2+
from unittest.mock import Mock, PropertyMock
3+
4+
import numpy as np
5+
import pytest
6+
7+
from encord.objects import LabelRowV2, Object, ObjectInstance, Shape
8+
from encord.objects.bitmask import BitmaskCoordinates
9+
from encord.orm.label_row import AnnotationTaskStatus, LabelRowMetadata, LabelStatus
10+
11+
bitmask_object = Object(
12+
uid=1, name="Mask", color="#D33115", shape=Shape.BITMASK, feature_node_hash="bitmask123", attributes=[]
13+
)
14+
get_child_by_hash = PropertyMock(return_value=bitmask_object)
15+
ontology_structure = Mock(get_child_by_hash=get_child_by_hash)
16+
ontology = Mock(structure=ontology_structure)
17+
18+
19+
def test_image_bitmask_dimension_validation():
20+
# Create label row metadata with specific dimensions (512x512)
21+
metadata = LabelRowMetadata(
22+
label_hash="test_label",
23+
branch_name="main",
24+
created_at=datetime.now(),
25+
last_edited_at=datetime.now(),
26+
data_hash="test_data",
27+
data_title="Test Image",
28+
data_type="IMAGE",
29+
data_link="",
30+
dataset_hash="test_dataset",
31+
dataset_title="Test Dataset",
32+
label_status=LabelStatus.NOT_LABELLED,
33+
annotation_task_status=AnnotationTaskStatus.QUEUED,
34+
workflow_graph_node=None,
35+
is_shadow_data=False,
36+
duration=None,
37+
frames_per_second=None,
38+
number_of_frames=1,
39+
height=512,
40+
width=512,
41+
audio_codec=None,
42+
audio_sample_rate=None,
43+
audio_num_channels=None,
44+
audio_bit_depth=None,
45+
)
46+
47+
# Create empty labels dict
48+
empty_labels = {
49+
"label_hash": "test_label",
50+
"branch_name": "main",
51+
"created_at": "Thu, 09 Feb 2023 14:12:03 UTC",
52+
"last_edited_at": "Thu, 09 Feb 2023 14:12:03 UTC",
53+
"data_hash": "test_data",
54+
"annotation_task_status": "QUEUED",
55+
"is_shadow_data": False,
56+
"dataset_hash": "test_dataset",
57+
"dataset_title": "Test Dataset",
58+
"data_title": "Test Image",
59+
"data_type": "image",
60+
"data_units": {
61+
"test_data": {
62+
"data_hash": "test_data",
63+
"data_title": "Test Image",
64+
"data_link": "",
65+
"data_type": "image/png",
66+
"data_sequence": "0",
67+
"width": 512,
68+
"height": 512,
69+
"labels": {"objects": [], "classifications": []},
70+
}
71+
},
72+
"object_answers": {},
73+
"classification_answers": {},
74+
"object_actions": {},
75+
"label_status": "LABEL_IN_PROGRESS",
76+
}
77+
78+
# Correct dimensions (512x512) should succeed
79+
label_row = LabelRowV2(metadata, Mock(), ontology)
80+
label_row.from_labels_dict(empty_labels)
81+
82+
correct_bitmask_coords = BitmaskCoordinates(np.zeros((512, 512), dtype=bool))
83+
correct_obj_instance = ObjectInstance(bitmask_object)
84+
correct_obj_instance.set_for_frames(coordinates=correct_bitmask_coords, frames=0)
85+
label_row.add_object_instance(correct_obj_instance)
86+
assert len(label_row.get_object_instances()) == 1
87+
label_row.to_encord_dict() # Serialization should also succeed
88+
89+
# Incorrect dimensions (256x256) should raise ValueError
90+
incorrect_bitmask_coords = BitmaskCoordinates(np.zeros((256, 256), dtype=bool))
91+
incorrect_obj_instance = ObjectInstance(bitmask_object)
92+
incorrect_obj_instance.set_for_frames(coordinates=incorrect_bitmask_coords, frames=0)
93+
label_row.add_object_instance(incorrect_obj_instance)
94+
assert len(label_row.get_object_instances()) == 2
95+
96+
with pytest.raises(ValueError, match="Bitmask dimensions don't match the media dimensions"):
97+
label_row.to_encord_dict()
98+
99+
100+
def test_image_group_bitmask_dimension_validation():
101+
metadata = LabelRowMetadata(
102+
label_hash="test_label",
103+
branch_name="main",
104+
created_at=datetime.now(),
105+
last_edited_at=datetime.now(),
106+
data_hash="test_data",
107+
data_title="Test Image Group",
108+
data_type="IMG_GROUP",
109+
data_link="",
110+
dataset_hash="test_dataset",
111+
dataset_title="Test Dataset",
112+
label_status=LabelStatus.NOT_LABELLED,
113+
annotation_task_status=AnnotationTaskStatus.QUEUED,
114+
workflow_graph_node=None,
115+
is_shadow_data=False,
116+
duration=None,
117+
frames_per_second=None,
118+
number_of_frames=2,
119+
height=None,
120+
width=None,
121+
audio_codec=None,
122+
audio_sample_rate=None,
123+
audio_num_channels=None,
124+
audio_bit_depth=None,
125+
)
126+
127+
# Create empty labels dict for image group with different dimensions per frame
128+
empty_labels = {
129+
"label_hash": "test_label",
130+
"branch_name": "main",
131+
"created_at": "Thu, 09 Feb 2023 14:12:03 UTC",
132+
"last_edited_at": "Thu, 09 Feb 2023 14:12:03 UTC",
133+
"data_hash": "test_data",
134+
"annotation_task_status": "QUEUED",
135+
"is_shadow_data": False,
136+
"dataset_hash": "test_dataset",
137+
"dataset_title": "Test Dataset",
138+
"data_title": "Test Image Group",
139+
"data_type": "img_group",
140+
"data_units": {
141+
"frame_0_hash": {
142+
"data_hash": "frame_0_hash",
143+
"data_title": "Frame 0",
144+
"data_link": "",
145+
"data_type": "image/png",
146+
"data_sequence": "0",
147+
"width": 512,
148+
"height": 512,
149+
"labels": {"objects": [], "classifications": []},
150+
},
151+
"frame_1_hash": {
152+
"data_hash": "frame_1_hash",
153+
"data_title": "Frame 1",
154+
"data_link": "",
155+
"data_type": "image/png",
156+
"data_sequence": "1",
157+
"width": 1024,
158+
"height": 768,
159+
"labels": {"objects": [], "classifications": []},
160+
},
161+
},
162+
"object_answers": {},
163+
"classification_answers": {},
164+
"object_actions": {},
165+
"label_status": "LABEL_IN_PROGRESS",
166+
}
167+
168+
label_row = LabelRowV2(metadata, Mock(), ontology)
169+
label_row.from_labels_dict(empty_labels)
170+
171+
# Add bitmask with correct dimensions for frame 0 (512x512)
172+
frame_0_correct_mask = BitmaskCoordinates(np.zeros((512, 512), dtype=bool))
173+
frame_0_correct_instance = ObjectInstance(bitmask_object)
174+
frame_0_correct_instance.set_for_frames(coordinates=frame_0_correct_mask, frames=0)
175+
label_row.add_object_instance(frame_0_correct_instance)
176+
177+
# Add bitmask with correct dimensions for frame 1 (1024x768)
178+
frame_1_correct_mask = BitmaskCoordinates(np.zeros((768, 1024), dtype=bool))
179+
frame_1_correct_instance = ObjectInstance(bitmask_object)
180+
frame_1_correct_instance.set_for_frames(coordinates=frame_1_correct_mask, frames=1)
181+
label_row.add_object_instance(frame_1_correct_instance)
182+
assert len(label_row.get_object_instances()) == 2
183+
label_row.to_encord_dict() # Both correct dimensions should serialize successfully
184+
185+
# Add bitmask with incorrect dimensions for frame 0 (256x256 instead of 512x512)
186+
frame_0_incorrect_mask = BitmaskCoordinates(np.zeros((256, 256), dtype=bool))
187+
frame_0_incorrect_instance = ObjectInstance(bitmask_object)
188+
frame_0_incorrect_instance.set_for_frames(coordinates=frame_0_incorrect_mask, frames=0, overwrite=True)
189+
label_row.add_object_instance(frame_0_incorrect_instance, force=True)
190+
191+
# Should fail on serialization
192+
with pytest.raises(ValueError, match="Bitmask dimensions don't match the media dimensions"):
193+
label_row.to_encord_dict()
194+
195+
# Fix frame 0, then add incorrect dimensions for frame 1
196+
label_row.remove_object(frame_0_incorrect_instance)
197+
label_row.to_encord_dict() # back to a valid serializable state
198+
199+
frame_1_incorrect_mask = BitmaskCoordinates(np.zeros((512, 512), dtype=bool))
200+
frame_1_incorrect_instance = ObjectInstance(bitmask_object)
201+
frame_1_incorrect_instance.set_for_frames(coordinates=frame_1_incorrect_mask, frames=1, overwrite=True)
202+
label_row.add_object_instance(frame_1_incorrect_instance, force=True)
203+
204+
# Should fail on serialization due to frame 1
205+
with pytest.raises(ValueError, match="Bitmask dimensions don't match the media dimensions"):
206+
label_row.to_encord_dict()
207+
208+
209+
def test_dicom_bitmask_dimension_validation():
210+
# Create label row metadata for DICOM with series-level dimensions (512x512)
211+
metadata = LabelRowMetadata(
212+
label_hash="test_label",
213+
branch_name="main",
214+
created_at=datetime.now(),
215+
last_edited_at=datetime.now(),
216+
data_hash="test_dicom_hash",
217+
data_title="Test DICOM",
218+
data_type="DICOM",
219+
data_link="",
220+
dataset_hash="test_dataset",
221+
dataset_title="Test Dataset",
222+
label_status=LabelStatus.NOT_LABELLED,
223+
annotation_task_status=AnnotationTaskStatus.QUEUED,
224+
workflow_graph_node=None,
225+
is_shadow_data=False,
226+
duration=None,
227+
frames_per_second=None,
228+
number_of_frames=2,
229+
height=512,
230+
width=512,
231+
audio_codec=None,
232+
audio_sample_rate=None,
233+
audio_num_channels=None,
234+
audio_bit_depth=None,
235+
)
236+
237+
dicom_labels = {
238+
"label_hash": "test_label",
239+
"branch_name": "main",
240+
"created_at": "Thu, 09 Feb 2023 14:12:03 UTC",
241+
"last_edited_at": "Thu, 09 Feb 2023 14:12:03 UTC",
242+
"data_hash": "test_dicom_hash",
243+
"annotation_task_status": "QUEUED",
244+
"is_shadow_data": False,
245+
"dataset_hash": "test_dataset",
246+
"dataset_title": "Test Dataset",
247+
"data_title": "Test DICOM",
248+
"data_type": "dicom",
249+
"data_units": {
250+
"test_dicom_hash": {
251+
"data_hash": "test_dicom_hash",
252+
"data_title": "Test DICOM",
253+
"data_type": "application/dicom",
254+
"data_sequence": 0,
255+
"labels": {
256+
"0": {
257+
"objects": [],
258+
"classifications": [],
259+
"metadata": {
260+
"dicom_instance_uid": "1.2.3.4.5.6.7.8.9.0",
261+
"multiframe_frame_number": None,
262+
"file_uri": "test/slice_0",
263+
"width": 512,
264+
"height": 512,
265+
},
266+
},
267+
"1": {
268+
"objects": [],
269+
"classifications": [],
270+
"metadata": {
271+
"dicom_instance_uid": "1.2.3.4.5.6.7.8.9.1",
272+
"multiframe_frame_number": None,
273+
"file_uri": "test/slice_1",
274+
"width": 10,
275+
"height": 10,
276+
},
277+
},
278+
},
279+
"metadata": {
280+
"patient_id": "test_patient",
281+
"study_uid": "1.2.3.4.5",
282+
"series_uid": "1.2.3.4.6",
283+
},
284+
"data_links": ["test/slice_0", "test/slice_1"],
285+
"width": 512,
286+
"height": 512,
287+
}
288+
},
289+
"object_answers": {},
290+
"classification_answers": {},
291+
"object_actions": {},
292+
"label_status": "LABEL_IN_PROGRESS",
293+
}
294+
295+
label_row = LabelRowV2(metadata, Mock(), ontology)
296+
label_row.from_labels_dict(dicom_labels)
297+
298+
slice_0_metadata = label_row.get_frame_view(0).metadata
299+
assert slice_0_metadata is not None
300+
assert slice_0_metadata.model_dump() == {
301+
"width": 512,
302+
"height": 512,
303+
"dicom_instance_uid": "1.2.3.4.5.6.7.8.9.0",
304+
"multiframe_frame_number": None,
305+
"file_uri": "test/slice_0",
306+
}
307+
308+
slice_1_metadata = label_row.get_frame_view(1).metadata
309+
assert slice_1_metadata is not None
310+
assert slice_1_metadata.model_dump() == {
311+
"width": 10,
312+
"height": 10,
313+
"dicom_instance_uid": "1.2.3.4.5.6.7.8.9.1",
314+
"multiframe_frame_number": None,
315+
"file_uri": "test/slice_1",
316+
}
317+
318+
# Add bitmask with correct dimensions for slice 0
319+
slice_0_correct_mask = BitmaskCoordinates(np.zeros((512, 512), dtype=bool))
320+
slice_0_correct_instance = ObjectInstance(bitmask_object)
321+
slice_0_correct_instance.set_for_frames(coordinates=slice_0_correct_mask, frames=0)
322+
label_row.add_object_instance(slice_0_correct_instance)
323+
assert len(label_row.get_object_instances()) == 1
324+
325+
# Add bitmask with correct dimensions for slice 1
326+
slice_1_correct_mask = BitmaskCoordinates(np.zeros((10, 10), dtype=bool))
327+
slice_1_correct_instance = ObjectInstance(bitmask_object)
328+
slice_1_correct_instance.set_for_frames(coordinates=slice_1_correct_mask, frames=1)
329+
label_row.add_object_instance(slice_1_correct_instance)
330+
assert len(label_row.get_object_instances()) == 2
331+
332+
# Both correct dimensions should serialize successfully
333+
label_row.to_encord_dict()
334+
335+
# Add bitmask with incorrect dimensions for slice 0 (256x256 instead of 512x512)
336+
slice_0_incorrect_mask = BitmaskCoordinates(np.zeros((256, 256), dtype=bool))
337+
slice_0_incorrect_instance = ObjectInstance(bitmask_object)
338+
slice_0_incorrect_instance.set_for_frames(coordinates=slice_0_incorrect_mask, frames=0, overwrite=True)
339+
label_row.add_object_instance(slice_0_incorrect_instance, force=True)
340+
341+
# Should fail on serialization
342+
with pytest.raises(ValueError, match="Bitmask dimensions don't match the media dimensions"):
343+
label_row.to_encord_dict()

0 commit comments

Comments
 (0)