feat: ED-1855 Dicom and Nifti Space

encord/constants/enums.py

@@ -71,3 +71,7 @@ class SpaceType(StringEnum):
     AUDIO = "audio"
     TEXT = "text"
     HTML = "html"
+    MEDICAL_FILE = "medical_file"
+    MEDICAL_STACK = "medical_stack"
+    SCENE_IMAGE = "scene_image"
+    POINT_CLOUD = "point_cloud"

encord/objects/ontology_labels_impl.py

@@ -106,10 +106,12 @@
 from encord.objects.spaces.base_space import Space, SpaceT
 from encord.objects.spaces.html_space import HTMLSpace
 from encord.objects.spaces.image_space import ImageSpace
+from encord.objects.spaces.multiframe_space.medical_file_space import MedicalFileSpace
+from encord.objects.spaces.multiframe_space.medical_stack_space import MedicalStackSpace
+from encord.objects.spaces.multiframe_space.video_space import VideoSpace
 from encord.objects.spaces.range_space.audio_space import AudioSpace
 from encord.objects.spaces.range_space.text_space import TextSpace
-from encord.objects.spaces.types import SpaceInfo
-from encord.objects.spaces.video_space import VideoSpace
+from encord.objects.spaces.types import ChildInfo, MedicalStackSpaceInfo, SpaceInfo
 from encord.objects.types import (
     AttributeDict,
     BaseFrameObject,
@@ -146,13 +148,16 @@
     "For this operation you will need to initialise labelling first. Call the `.initialise_labels()` to do so first."
 )
 
+
 # Type mapping for runtime validation in get_space
 _SPACE_TYPE_TO_CLASS = {
     "video": VideoSpace,
     "image": ImageSpace,
     "audio": AudioSpace,
     "text": TextSpace,
     "html": HTMLSpace,
+    "medical_file": MedicalFileSpace,
+    "medical_stack": MedicalStackSpace,
 }
 
 
@@ -811,6 +816,18 @@ def _get_space(self, *, id: str, type_: Literal["audio"]) -> AudioSpace:
     def _get_space(self, *, id: str, type_: Literal["text"]) -> TextSpace:
         pass
 
+    @overload
+    def _get_space(self, *, id: str, type_: Literal["html"]) -> HTMLSpace:
+        pass
+
+    @overload
+    def _get_space(self, *, id: str, type_: Literal["medical_file"]) -> MedicalFileSpace:
+        pass
+
+    @overload
+    def _get_space(self, *, id: str, type_: Literal["medical_stack"]) -> MedicalStackSpace:
+        pass
+
     @overload
     def _get_space(self, *, layout_key: str, type_: Literal["video"]) -> VideoSpace:
         pass
@@ -828,19 +845,23 @@ def _get_space(self, *, layout_key: str, type_: Literal["text"]) -> TextSpace:
         pass
 
     @overload
-    def _get_space(self, *, id: str, type_: Literal["html"]) -> HTMLSpace:
+    def _get_space(self, *, layout_key: str, type_: Literal["html"]) -> HTMLSpace:
         pass
 
     @overload
-    def _get_space(self, *, layout_key: str, type_: Literal["html"]) -> HTMLSpace:
+    def _get_space(self, *, layout_key: str, type_: Literal["medical_file"]) -> MedicalFileSpace:
+        pass
+
+    @overload
+    def _get_space(self, *, layout_key: str, type_: Literal["medical_stack"]) -> MedicalStackSpace:
         pass
 
     def _get_space(
         self,
         *,
         id: Optional[str] = None,
         layout_key: Optional[str] = None,
-        type_: Literal["video", "image", "audio", "text", "html"],
+        type_: Literal["video", "image", "audio", "text", "html", "medical_file", "medical_stack"],
     ) -> Space:
         """Retrieves a single space which matches the specified id and type.
 
@@ -2470,7 +2491,7 @@ def _initiate_spaces(
                 continue
 
             # Store layout_key -> space_id mapping if child_info is present
-            child_info = space_info.get("child_info")
+            child_info = cast(Optional[ChildInfo], space_info.get("child_info", None))
             if child_info is not None:
                 self._layout_key_to_space_id[child_info["layout_key"]] = space_id
 
@@ -2510,6 +2531,23 @@ def _initiate_spaces(
                     label_row=self,
                 )
                 res[space_id] = html_space
+            elif space_info["space_type"] == SpaceType.MEDICAL_FILE:
+                medical_file_space = MedicalFileSpace(
+                    space_id=space_id,
+                    label_row=self,
+                    number_of_frames=space_info["number_of_frames"],
+                    width=space_info["width"],
+                    height=space_info["height"],
+                )
+                res[space_id] = medical_file_space
+            elif space_info["space_type"] == SpaceType.MEDICAL_STACK:
+                medical_stack_space = MedicalStackSpace(space_id=space_id, label_row=self, frames=space_info["frames"])
+                res[space_id] = medical_stack_space
+            elif space_info["space_type"] == SpaceType.SCENE_IMAGE or space_info["space_type"] == SpaceType.POINT_CLOUD:
+                # TODO: Implement Scene Images
+                pass
+            else:
+                exhaustive_guard(space_info["space_type"], message="Missing initialisation for space.")
 
         return res
 
@@ -2545,6 +2583,21 @@ def _parse_space_labels(
                 html_space._parse_space_dict(
                     space_info, object_answers=object_answers, classification_answers=classification_answers
                 )
+            elif space_info["space_type"] == SpaceType.MEDICAL_FILE:
+                medical_file_space = self._get_space(id=space_id, type_="medical_file")
+                medical_file_space._parse_space_dict(
+                    space_info, object_answers=object_answers, classification_answers=classification_answers
+                )
+            elif space_info["space_type"] == SpaceType.MEDICAL_STACK:
+                medical_stack_space = self._get_space(id=space_id, type_="medical_stack")
+                medical_stack_space._parse_space_dict(
+                    space_info, object_answers=object_answers, classification_answers=classification_answers
+                )
+            elif space_info["space_type"] == SpaceType.SCENE_IMAGE or space_info["space_type"] == SpaceType.POINT_CLOUD:
+                # TODO: Enable this when we implement Scene images
+                pass
+            else:
+                exhaustive_guard(space_info["space_type"], message="Missing implementation for parsing space labels.")
 
     def _parse_label_row_metadata(self, label_row_metadata: LabelRowMetadata) -> LabelRowV2.LabelRowReadOnlyData:
         data_type = DataType.from_upper_case_string(label_row_metadata.data_type)

encord/objects/spaces/annotation/geometric_annotation.py

@@ -14,7 +14,7 @@
 if TYPE_CHECKING:
     from encord.objects import ClassificationInstance, ObjectInstance
     from encord.objects.spaces.image_space import ImageSpace
-    from encord.objects.spaces.video_space import VideoSpace
+    from encord.objects.spaces.multiframe_space.multiframe_space import MultiFrameSpace
 
 
 @dataclass
@@ -62,13 +62,13 @@ def _check_if_annotation_is_valid(self) -> None:
 class _GeometricFrameObjectAnnotation(_ObjectAnnotation):
     """Annotations for multi-frame geometric object labels (e.g. Video)."""
 
-    def __init__(self, space: VideoSpace, object_instance: ObjectInstance, frame: int):
+    def __init__(self, space: MultiFrameSpace, object_instance: ObjectInstance, frame: int):
         super().__init__(space, object_instance)
-        self._space: VideoSpace = space
+        self._space: MultiFrameSpace = space
         self._frame = frame
 
     @property
-    def space(self) -> VideoSpace:
+    def space(self) -> MultiFrameSpace:
         return self._space
 
     @property
@@ -107,13 +107,13 @@ def _check_if_annotation_is_valid(self) -> None:
 class _FrameClassificationAnnotation(_ClassificationAnnotation):
     """Annotations for multi-frame classifications (e.g. Video)."""
 
-    def __init__(self, space: VideoSpace, classification_instance: ClassificationInstance, frame: int):
+    def __init__(self, space: MultiFrameSpace, classification_instance: ClassificationInstance, frame: int):
         super().__init__(space, classification_instance)
-        self._space: VideoSpace = space
+        self._space: MultiFrameSpace = space
         self._frame = frame
 
     @property
-    def space(self) -> VideoSpace:
+    def space(self) -> MultiFrameSpace:
         return self._space
 
     @property

encord/objects/spaces/image_space.py

@@ -22,8 +22,8 @@
 )
 from encord.objects.spaces.annotation.global_annotation import _GlobalClassificationAnnotation
 from encord.objects.spaces.base_space import Space
+from encord.objects.spaces.multiframe_space.multiframe_space import FrameOverlapStrategy
 from encord.objects.spaces.types import ImageSpaceInfo, SpaceInfo
-from encord.objects.spaces.video_space import FrameOverlapStrategy
 from encord.objects.types import (
     AttributeDict,
     ClassificationAnswer,
@@ -40,7 +40,6 @@
 if TYPE_CHECKING:
     from encord.objects import ClassificationInstance, ObjectInstance
     from encord.objects.ontology_labels_impl import LabelRowV2
-    from encord.objects.ontology_object import Object
 
 
 class ImageSpace(Space[_GeometricObjectAnnotation, _GlobalClassificationAnnotation, FrameOverlapStrategy]):

encord/objects/spaces/multiframe_space/medical_file_space.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from encord.constants.enums import SpaceType
+from encord.objects.spaces.multiframe_space.multiframe_space import MultiFrameSpace
+from encord.objects.spaces.types import MedicalFileSpaceInfo
+from encord.objects.types import LabelBlob
+
+if TYPE_CHECKING:
+    from encord.objects.ontology_labels_impl import LabelRowV2
+
+logger = logging.getLogger(__name__)
+
+
+class MedicalFileSpace(MultiFrameSpace):
+    """Space for medical files (e.g. DICOM, NIfTI)."""
+
+    def __init__(
+        self,
+        space_id: str,
+        label_row: LabelRowV2,
+        number_of_frames: int,
+        width: int,
+        height: int,
+    ):
+        super().__init__(space_id, label_row, number_of_frames=number_of_frames)
+        self._number_of_frames = number_of_frames
+        self._width = width
+        self._height = height
+
+    def _get_frame_dimensions(self, frame: int) -> tuple[int, int]:
+        return self._width, self._height
+
+    def _to_space_dict(self) -> MedicalFileSpaceInfo:
+        """Export space to dictionary format."""
+        frame_labels = self._build_frame_labels_dict()
+
+        return MedicalFileSpaceInfo(
+            space_type=SpaceType.MEDICAL_FILE,
+            labels=frame_labels,
+            number_of_frames=self._number_of_frames,
+            width=self._width,
+            height=self._height,
+        )

encord/objects/spaces/multiframe_space/medical_stack_space.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, List, Optional
+
+from encord.constants.enums import SpaceType
+from encord.exceptions import LabelRowError
+from encord.objects.spaces.multiframe_space.multiframe_space import MultiFrameSpace
+from encord.objects.spaces.types import DicomFrameInfo, MedicalStackSpaceInfo
+from encord.objects.types import LabelBlob
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from encord.objects.ontology_labels_impl import LabelRowV2
+
+
+class MedicalStackSpace(MultiFrameSpace):
+    """Space for medical files (e.g. DICOM, NIfTI)."""
+
+    def __init__(self, space_id: str, label_row: LabelRowV2, frames: List[DicomFrameInfo]):
+        super().__init__(space_id, label_row, number_of_frames=len(frames))
+        self._frames = frames
+
+    def _get_frame_dimensions(self, frame_number: int) -> tuple[int, int]:
+        if frame_number >= self._number_of_frames:
+            raise LabelRowError(
+                f"Invalid frame number '{frame_number}'. This file only has {self._number_of_frames} frames."
+            )
+
+        frame = self._frames[frame_number]
+        return frame["width"], frame["height"]
+
+    def _to_space_dict(self) -> MedicalStackSpaceInfo:
+        """Export space to dictionary format."""
+        frame_labels = self._build_frame_labels_dict()
+
+        return MedicalStackSpaceInfo(space_type=SpaceType.MEDICAL_STACK, labels=frame_labels, frames=self._frames)