retuve.classes.seg

  1# Copyright 2024 Adam McArthur
  2#
  3# Licensed under the Apache License, Version 2.0 (the "License");
  4# you may not use this file except in compliance with the License.
  5# You may obtain a copy of the License at
  6#
  7#     http://www.apache.org/licenses/LICENSE-2.0
  8#
  9# Unless required by applicable law or agreed to in writing, software
 10# distributed under the License is distributed on an "AS IS" BASIS,
 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12# See the License for the specific language governing permissions and
 13# limitations under the License.
 14
 15"""
 16Segmentaition Classes
 17
 18These are the classes you need to store your model results in.
 19
 20They are used so Retuve can understand the results of your model.
 21
 22Below is an example of how to write a custom AI model and config for Retuve.
 23
 24https://github.com/radoss-org/retuve/tree/main/examples/ai_plugins/custom_ai_and_config.py
 25
 26```python
 27.. include:: ../../examples/ai_plugins/custom_ai_and_config.py
 28```
 29
 30This file can then be used with the UI and Trak using.
 31
 32```bash
 33retuve --task trak --keyphrase_file config.py
 34```
 35"""
 36
 37from typing import Annotated, Iterable, List, Literal, Tuple, TypeVar
 38
 39import numpy as np
 40import numpy.typing as npt
 41from PIL import Image
 42from radstract.data.colors import get_unique_colours
 43
 44from retuve.hip_us.classes.enums import HipLabelsUS
 45from retuve.hip_xray.classes import HipLabelsXray
 46
 47DType = TypeVar("DType", bound=np.generic)
 48
 49NDArrayImg_NxNx3 = Annotated[npt.NDArray[DType], Literal["N", "N", 3]]
 50NDArrayImg_NxNx3_AllWhite = Annotated[npt.NDArray[DType], Literal["N", "N", 3]]
 51MidLine = List[Tuple[int, int]]
 52
 53
 54class SegObject:
 55    """
 56    Class for holding a single segmentation object.
 57    """
 58
 59    def __init__(
 60        self,
 61        points: List[Tuple[int, int]] = None,
 62        clss: HipLabelsUS = None,
 63        mask: NDArrayImg_NxNx3_AllWhite = None,
 64        conf: float = None,
 65        box: Tuple[int, int, int, int] = None,
 66        empty: bool = False,
 67    ):
 68        """
 69        :param points: List of points that make up the object.
 70        :param clss: Class of the object.
 71        :param mask: Mask of the object.
 72        :param conf: Confidence of the object.
 73        :param box: Bounding box of the object.
 74        :param empty: Is the object empty.
 75
 76        :attr midline: Midline of the object (Only for Hip Ultrasound).
 77        :attr midline_moved: Midline of the object after it
 78              has been moved (Only for Hip Ultrasound).
 79
 80        :raises ValueError: If any of the parameters are not as expected.
 81        """
 82        self.points = points
 83        self.cls = clss
 84        self.mask = mask
 85        self.box = box
 86        self.conf = conf
 87        self.empty = empty
 88
 89        # Only for hip_us
 90        self.midline: MidLine = None
 91        self.midline_moved: MidLine = None
 92
 93        if empty:
 94            return
 95
 96        if points is not None:
 97            for point in points:
 98                if len(point) != 2:
 99                    raise ValueError("Point is not a tuple of length 2")
100
101        if type(mask) != np.ndarray:
102            raise ValueError("mask is not a numpy array")
103
104        # check mask shape
105        if len(mask.shape) != 3 or mask.shape[2] != 3:
106            raise ValueError("mask is not RGB as required")
107
108        colors = get_unique_colours(array=mask)
109        # check each pixel in mask is either rgb white or black
110        if (
111            any(color not in [(0, 0, 0), (255, 255, 255)] for color in colors)
112            or len(colors) > 2
113        ):
114            raise ValueError("mask is not all white")
115
116        if box is not None:
117            if len(box) != 4:
118                raise ValueError("box is not a tuple of length 4")
119
120        if clss is None or not (
121            isinstance(clss, int)
122            or isinstance(clss, HipLabelsUS)
123            or isinstance(clss, HipLabelsXray)
124        ):
125            raise ValueError("clss is None or wrong type")
126
127        if conf is not None and not 0 <= conf <= 1:
128            raise ValueError("conf is not None and not between 0 and 1")
129
130    def __str__(self):
131        return f"SegObject({self.cls}, {self.conf}, {self.points})"
132
133    def area(self):
134        """
135        Returns the area of the object.
136        """
137        # Use the mask to calculate the area
138        if self.mask is None:
139            return 0
140        return np.sum(self.mask[:, :, 0] == 255)
141
142    def flip_horizontally(self, img_width: int):
143        """
144        Flips the object horizontally.
145
146        :param img_width: Width of the image.
147        """
148        if self.empty:
149            return
150
151        if self.box is not None:
152            self.box = (
153                img_width - self.box[2],
154                img_width - self.box[0],
155                self.box[1],
156                self.box[3],
157            )
158
159        if self.points is not None:
160            self.points = [(img_width - x, y) for x, y in self.points]
161
162        if self.midline is not None:
163            self.midline = np.array(
164                [(y, img_width - x) for y, x in self.midline]
165            )
166
167        if self.midline_moved is not None:
168            self.midline_moved = np.array(
169                [(y, img_width - x) for y, x in self.midline_moved]
170            )
171
172        if self.mask is not None:
173            self.mask = np.flip(self.mask, axis=1)
174
175
176class SegFrameObjects:
177    """
178    Class for holding a frame of segmentation objects.
179    """
180
181    def __init__(
182        self, img: NDArrayImg_NxNx3, seg_objects: list[SegObject] = None
183    ):
184        """
185        :param img: Image of the frame.
186        :param seg_objects: List of segmentation objects in the frame.
187        """
188        if type(img) != np.ndarray:
189            raise ValueError("img is not a numpy array")
190
191        if seg_objects is None:
192            seg_objects = []
193
194        self.seg_objects = seg_objects
195        self.img = img
196
197    def __iter__(self) -> Iterable[SegObject]:
198        return iter(self.seg_objects)
199
200    def append(self, seg_result: SegObject):
201        self.seg_objects.append(seg_result)
202
203    def __getitem__(self, index):
204        return self.seg_objects[index]
205
206    def __setitem__(self, index, value):
207        self.seg_objects[index] = value
208
209    def __len__(self):
210        return len(self.seg_objects)
211
212    @classmethod
213    def empty(cls, img: Image.Image):
214        """
215        Returns an empty SegFrameObjects object.
216
217        :param img: Image of the frame.
218        """
219
220        return cls(img=img, seg_objects=[SegObject(empty=True)])
221
222    def __str__(self):
223        return f"SegFrameObjects({self.seg_objects})"