retuve.classes.draw

  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"""
 16Class for creating visual overlays on images.
 17"""
 18
 19from enum import Enum
 20from typing import List, Literal, Tuple
 21
 22import numpy as np
 23from PIL import Image, ImageDraw
 24from radstract.data.colors import LabelColours
 25
 26from retuve.classes.seg import SegFrameObjects
 27from retuve.keyphrases.config import Config
 28
 29
 30class DrawTypes(Enum):
 31    """
 32    Enum for the different types of overlays that can be drawn.
 33
 34    """
 35
 36    LINES = "lines"
 37    SEGS = "segs"
 38    TEXT = "text"
 39    POINTS = "points"
 40    CIRCLE = "circle"
 41    RECTANGLE = "rectangle"
 42
 43    @classmethod
 44    def ALL(cls) -> List["DrawTypes"]:
 45        """
 46        Specified in order of drawing.
 47        """
 48        return [
 49            cls.SEGS,
 50            cls.LINES,
 51            cls.CIRCLE,
 52            cls.RECTANGLE,
 53            cls.POINTS,
 54            cls.TEXT,
 55        ]
 56
 57    @classmethod
 58    def type_to_func(cls, draw, dtype):
 59        """
 60        Returns the function to use for a given overlay type.
 61
 62        :param draw: ImageDraw object.
 63        :param dtype: Type of overlay to draw.
 64
 65        :return: Function to use for a given overlay type.
 66        """
 67        if dtype == cls.LINES:
 68            return draw.line
 69        elif dtype == cls.SEGS:
 70            return draw.polygon
 71        elif dtype == cls.TEXT:
 72            return draw.text
 73        elif dtype == cls.POINTS:
 74            return draw.point
 75        elif dtype == cls.CIRCLE:
 76            return draw.ellipse
 77        elif dtype == cls.RECTANGLE:
 78            return draw.rectangle
 79
 80
 81class Overlay:
 82    """
 83    Class for creating visual overlays on images.
 84    """
 85
 86    def __init__(self, shape: Tuple[int, int], config: Config):
 87        """
 88        :param shape: Shape of the image to overlay on.
 89        :param config: Config object.
 90
 91        :attr shape: Shape of the image to overlay on.
 92        :attr overlays: Dictionary of overlays.
 93        :attr config: Config object.
 94        """
 95        self.shape = (shape[0], shape[1], 4)
 96        self.config = config
 97        self.operations = {}
 98
 99        for drtype in DrawTypes.ALL():
100            self.operations[drtype] = []
101
102    def add_operation(self, optype, *args, **kwargs):
103        # Store the drawing operation in the operations list
104        self.operations[optype].append(((args, kwargs)))
105
106    def apply_to_image(self, image):
107
108        if not any(self.operations.values()):
109            return image
110
111        image = Image.fromarray(image)
112        draw = ImageDraw.Draw(image, "RGBA")
113
114        # Execute all stored operations
115        # go in order of segs, lines, points, text
116        for optype in DrawTypes.ALL():
117            for args, kwargs in self.operations[optype]:
118                func = DrawTypes.type_to_func(draw, optype)
119                func(*args, **kwargs)
120
121        final_image = np.array(image, dtype=np.uint8)
122
123        # remove alpha channel
124        final_image = final_image[:, :, :3]
125
126        return final_image
127
128    def get_nifti_frame(
129        self, seg_frame_objs: SegFrameObjects, shape: list
130    ) -> Image.Image:
131        """
132        Create an image with segmentation overlay ready for NIFTI storage.
133
134        :param seg_frame_objs: Segmentation frame objects.
135        :param shape: Shape of the image.
136
137        :return: Image with segmentation overlay.
138        """
139
140        # Create empty frame of write shape
141        seg_overlay = Image.new("RGB", shape[:2], (0, 0, 0))
142        draw = ImageDraw.Draw(seg_overlay)
143
144        for seg_obj in seg_frame_objs:
145            if not seg_obj.empty:
146                draw.polygon(
147                    seg_obj.points,
148                    fill=LabelColours.get_color_from_index(
149                        seg_obj.cls.value + 1
150                    ),
151                )
152
153        return seg_overlay
154
155    def draw_cross(self, point: Tuple[int, int]):
156        """
157        Draws a cross of a given radius at a specified point on the overlay.
158
159        :param point: Point to draw the cross at.
160        """
161        x, y = point
162
163        radius = self.config.visuals.points_radius
164        color = self.config.visuals.points_color.rgba()
165
166        self.add_operation(
167            DrawTypes.LINES,
168            (x - radius, y, x + radius, y),
169            fill=color,
170            width=self.config.visuals.line_thickness,
171        )
172        self.add_operation(
173            DrawTypes.LINES,
174            (x, y - radius, x, y + radius),
175            fill=color,
176            width=self.config.visuals.line_thickness,
177        )
178
179    def draw_segmentation(self, points: List[Tuple[int, int]]):
180        """
181        Draws a polygon on the overlay.
182
183        :param points: List of points to draw the polygon.
184        """
185
186        self.add_operation(
187            DrawTypes.SEGS,
188            points,
189            fill=self.config.visuals.seg_color.rgba(
190                self.config.visuals.seg_alpha
191            ),
192        )
193
194    def draw_box(self, box: Tuple[int, int, int, int], grafs: bool = False):
195        """
196        Draws a bounding box on the overlay.
197
198        :param box: Bounding box to draw.
199        :param grafs: Whether to draw with the graf frame colours or not.
200
201        """
202        x1, y1, x2, y2 = box
203        start_point = (int(x1), int(y1))
204        end_point = (int(x2), int(y2))  # bottom right corner
205
206        if grafs:
207            color = self.config.visuals.graf_color.rgba()
208            width = 10
209        else:
210            color = self.config.visuals.bounding_box_color.rgba()
211            width = self.config.visuals.bounding_box_thickness
212
213        self.add_operation(
214            DrawTypes.RECTANGLE,
215            [start_point, end_point],
216            outline=color,
217            width=width,
218        )
219
220    def draw_skeleton(self, skel: List[Tuple[int, int]]):
221        """
222        Draws a skeleton on the overlay.
223
224        :param skel: List of points to draw the skeleton.
225        """
226
227        for point in skel:
228            y, x = point
229            self.add_operation(
230                DrawTypes.POINTS,
231                (x, y),
232                fill=self.config.hip.midline_color.rgba(),
233            )
234
235    def draw_lines(
236        self, line_points: List[Tuple[Tuple[int, int], Tuple[int, int]]]
237    ):
238        """
239        Draws lines on the overlay.
240
241        :param line_points: List of tuples of points to draw lines between.
242
243        """
244
245        for point1, point2 in line_points:
246            self.add_operation(
247                DrawTypes.LINES,
248                (tuple(point1), tuple(point2)),
249                fill=self.config.visuals.line_color.rgba(),
250                width=self.config.visuals.line_thickness,
251            )
252
253    def draw_text(
254        self,
255        label_text: str,
256        x1: int,
257        y1: int,
258        header: Literal["h1", "h2"] = "h1",
259        grafs: bool = False,
260    ):
261        """
262        Draws text on the overlay.
263
264        :param label_text: Text to draw.
265        :param x1: x coordinate to draw text.
266        :param y1: y coordinate to draw text.
267        :param header: Header type to use.
268        :param grafs: Whether to draw with the graf frame colours or not.
269
270        """
271        if header == "h1":
272            font = self.config.visuals.font_h1
273        elif header == "h2":
274            font = self.config.visuals.font_h2
275
276        temp_draw = ImageDraw.Draw(Image.new("RGBA", self.shape[:2]))
277
278        bbox = temp_draw.textbbox((x1, y1), label_text, font=font)
279
280        if grafs:
281            color = self.config.visuals.graf_color.rgba()
282        else:
283            color = self.config.visuals.background_text_color.rgba()
284
285        self.add_operation(
286            DrawTypes.RECTANGLE,
287            bbox,
288            fill=color,
289        )
290
291        self.add_operation(
292            DrawTypes.TEXT,
293            (x1, y1),
294            label_text,
295            font=font,
296            fill=self.config.visuals.text_color.rgba(),
297        )
298
299    def draw_circle(self, point: Tuple[int, int], radius: int):
300        """
301        Draws a circle on the overlay.
302
303        :param point: Center of the circle.
304        :param radius: Radius of the circle.
305        """
306
307        x, y = point
308
309        self.add_operation(
310            DrawTypes.CIRCLE,
311            (x - radius, y - radius, x + radius, y + radius),
312            outline=self.config.visuals.points_color.rgba(),
313            width=self.config.visuals.line_thickness,
314        )