retuve.hip_us.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"""
 16All unique drawing functions for Hip Ultrasound
 17"""
 18
 19import io
 20import textwrap
 21import time
 22from typing import List, Tuple
 23
 24import cv2
 25import numpy as np
 26import plotly.graph_objects as go
 27from attr import has
 28from PIL import Image, ImageOps
 29from radstract.data.nifti import NIFTI, convert_images_to_nifti_labels
 30
 31from retuve.classes.draw import Overlay
 32from retuve.classes.seg import SegFrameObjects
 33from retuve.draw import (TARGET_SIZE, draw_landmarks, draw_seg,
 34                         resize_data_for_display)
 35from retuve.hip_us.classes.enums import Side
 36from retuve.hip_us.classes.general import HipDatasUS, HipDataUS
 37from retuve.hip_us.handlers.side import get_side_metainfo
 38from retuve.hip_us.metrics.alpha import draw_alpha
 39from retuve.hip_us.metrics.coverage import draw_coverage
 40from retuve.hip_us.metrics.curvature import draw_curvature
 41from retuve.hip_us.multiframe import FemSphere
 42from retuve.hip_us.multiframe.models import circle_radius_at_z
 43from retuve.keyphrases.config import Config
 44from retuve.logs import log_timings, ulogger
 45
 46
 47def draw_fem_head(
 48    hip: HipDataUS, fem_sph: FemSphere, overlay: Overlay, z_gap: float
 49) -> Overlay:
 50    """
 51    Draw the femoral head on the image
 52
 53    :param hip: The Hip Datas
 54    :param fem_sph: The Femoral Sphere
 55    :param overlay: The Overlay to draw on
 56
 57    :return: The Drawn Overlay
 58    """
 59    # Get radius at z
 60    radius = circle_radius_at_z(
 61        fem_sph.radius, fem_sph.center[2], z_gap * hip.frame_no
 62    )
 63
 64    # draw the circle
 65    overlay.draw_circle((fem_sph.center[0], fem_sph.center[1]), radius)
 66
 67    return overlay
 68
 69
 70def draw_hips_us(
 71    hip_datas: HipDatasUS,
 72    results: List[SegFrameObjects],
 73    fem_sph: FemSphere,
 74    config: Config,
 75) -> tuple[List[np.ndarray], NIFTI]:
 76    """
 77    Draw the hip ultrasound images
 78
 79    :param hip_datas: The Hip Datas
 80    :param results: The Segmentation Results
 81    :param fem_sph: The Femoral Sphere
 82    :param config: The Config
 83
 84    :return: The Drawn Images
 85    """
 86    draw_timings = []
 87    image_arrays = []
 88    nifti_frames = []
 89    nifti = None
 90
 91    for i, (hip, seg_frame_objs) in enumerate(zip(hip_datas, results)):
 92        start = time.time()
 93
 94        final_hip, final_seg_frame_objs, final_image = resize_data_for_display(
 95            hip, seg_frame_objs
 96        )
 97
 98        overlay = Overlay(
 99            (final_image.shape[0], final_image.shape[1], 3), config
100        )
101
102        overlay = draw_seg(final_seg_frame_objs, overlay, config)
103
104        overlay = draw_landmarks(final_hip, overlay)
105
106        overlay = draw_alpha(final_hip, overlay, config)
107        overlay = draw_coverage(final_hip, overlay, config)
108        overlay = draw_curvature(final_hip, overlay, config)
109
110        if fem_sph and config.hip.display_fem_guess:
111            overlay = draw_fem_head(
112                final_hip,
113                fem_sph,
114                overlay,
115                config.hip.z_gap,
116            )
117
118        graf_conf = None
119        if hasattr(hip_datas, "graf_confs"):
120            graf_conf = hip_datas.graf_confs[hip.frame_no]
121
122        overlay, is_graf = draw_other(
123            final_hip,
124            final_seg_frame_objs,
125            hip_datas.graf_frame,
126            overlay,
127            final_image.shape[:2],
128            config,
129            graf_conf,
130        )
131
132        if config.hip.display_bad_frame_reasons and hasattr(
133            hip_datas, "bad_frame_reasons"
134        ):
135            if hip.frame_no in hip_datas.bad_frame_reasons:
136                overlay.draw_text(
137                    hip_datas.bad_frame_reasons[hip.frame_no],
138                    final_image.shape[1] // 2,
139                    final_image.shape[0] - 100,
140                    header="h2",
141                )
142
143        img = overlay.apply_to_image(final_image)
144
145        if config.seg_export:
146            original_image = seg_frame_objs.img
147            test = overlay.get_nifti_frame(
148                seg_frame_objs,
149                # NOTE(sharpz7) I do not know why this needs to be reversed
150                (original_image.shape[1], original_image.shape[0]),
151            )
152            nifti_frames.append(test)
153
154        # if its the graf frame, append 5 copies
155        repeats = len(results) // 6 if is_graf else 1
156        for _ in range(repeats):
157            image_arrays.append(img)
158
159        draw_timings.append(time.time() - start)
160
161    start = time.time()
162    if config.seg_export:
163        frames = nifti_frames
164
165        if "Swapped Post and Ant" in hip_datas.recorded_error.errors:
166            frames = nifti_frames[::-1]
167
168        # convert to nifti
169        nifti = convert_images_to_nifti_labels(frames)
170    else:
171        nifti = None
172
173    if (time.time() - start) > 0.1:
174        log_timings([time.time() - start], title="Nifti Conversion Speed:")
175
176    log_timings(draw_timings, title="Drawing Speed:")
177
178    return image_arrays, nifti
179
180
181def draw_other(
182    hip: HipDataUS,
183    seg_frame_objs: SegFrameObjects,
184    graf_frame: int,
185    overlay: Overlay,
186    shape: tuple,
187    config: Config,
188    graf_conf: float = None,
189) -> Tuple[Overlay, bool]:
190    """
191    Draw the other meta information on the image
192
193    :param hip: The Hip Data
194    :param seg_frame_objs: The Segmentation Frame Objects
195    :param graf_frame: The Graf Frame
196    :param overlay: The Overlay
197    :param config: The Config
198    :param shape: The Shape of the Image
199
200    :return: The Drawn Overlay
201    """
202
203    if config.hip.display_frame_no:
204        # Draw Frame No
205        overlay.draw_text(
206            f"Frame: {hip.frame_no}",
207            0,
208            0,
209            header="h1",
210        )
211
212    if graf_conf is not None and config.hip.display_graf_conf:
213        overlay.draw_text(
214            f"Graf Confidence: {graf_conf:.2f}",
215            shape[0] - 100,
216            0,
217            header="h1",
218        )
219
220    is_graf = hip.frame_no == graf_frame and hip.landmarks is not None
221
222    # Check if graf alpha
223    if is_graf:
224        overlay.draw_text(
225            f"Grafs Frame",
226            int(hip.landmarks.left[0]),
227            int(hip.landmarks.left[1] - 120),
228            header="h2",
229            grafs=True,
230        )
231
232        # add a border
233        overlay.draw_box(
234            (0, 0, shape[1], shape[0]),
235            grafs=True,
236        )
237
238    if hip.landmarks and hip.landmarks.right and config.hip.display_side:
239        xl, yl = hip.landmarks.right
240        overlay.draw_text(
241            f"side: {Side.get_name(hip.side)}",
242            int(xl),
243            int(yl),
244            header="h2",
245        )
246
247    if config.hip.draw_midline:
248        for seg_obj in seg_frame_objs:
249            if seg_obj.midline is not None:
250                overlay.draw_skeleton(seg_obj.midline)
251
252            if seg_obj.midline_moved is not None:
253                overlay.draw_skeleton(seg_obj.midline_moved)
254
255    if config.hip.draw_side_metainfo:
256        closest_illium, mid = get_side_metainfo(hip, seg_frame_objs)
257        # these are points, draw them
258        if closest_illium is not None:
259            overlay.draw_cross(closest_illium)
260            overlay.draw_cross(mid)
261
262    return overlay, is_graf
263
264
265def draw_table(shape: tuple, hip_datas: HipDatasUS) -> np.ndarray:
266    """
267    Draw the table of the metrics onto an image
268
269    :param shape: The Shape of the Image
270    :param hip_datas: The Hip Datas
271
272    :return: The Image with the Table and any errors
273    """
274    start = time.time()
275
276    # Create empty image with the specified shape
277    empty_img = np.zeros((shape[1], shape[0], 3), dtype=np.uint8)
278
279    # Find new shape by running 1024 algo
280    shape = ImageOps.contain(Image.fromarray(empty_img), (TARGET_SIZE)).size[
281        :2
282    ]
283
284    headers = [""] + hip_datas.metrics[0].names()
285    values = []
286
287    for metrics in reversed(hip_datas.metrics):
288        values.append(metrics.dump())
289
290    # Rotate and transpose the values list
291    values = list(zip(*values[::-1]))
292
293    # Create a Plotly figure for the table
294    fig = go.Figure(
295        data=[
296            go.Table(
297                header=dict(values=headers, font=dict(size=18), height=50),
298                cells=dict(
299                    values=values,
300                    fill=dict(color=["paleturquoise", "white"]),
301                    font=dict(size=18),
302                    height=50,
303                ),
304            )
305        ]
306    )
307
308    MARGIN = 30
309    fig.update_layout(
310        autosize=False,
311        width=shape[1],
312        height=shape[0],
313        margin=dict(l=MARGIN, r=MARGIN, b=MARGIN, t=MARGIN),
314    )
315
316    # Save the figure as bytes
317    img_bytes = fig.to_image(format="png", width=shape[1], height=shape[0])
318
319    # Convert image bytes to numpy array
320    data_image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
321
322    # Text Wrapping Logic
323    recorded_error_text = str(hip_datas.recorded_error)
324
325    # Set the maximum width for the text box in pixels
326    max_text_width = data_image.shape[1] - 20  # Margin of 10px from both sides
327    font_scale = 0.8
328    font = cv2.FONT_HERSHEY_SIMPLEX
329    font_thickness = 2
330
331    # Estimate the width of each character based on the font
332    (text_width, text_height), _ = cv2.getTextSize(
333        recorded_error_text, font, font_scale, font_thickness
334    )
335    char_width = (
336        text_width // len(recorded_error_text) if recorded_error_text else 1
337    )  # Estimate width of one char
338
339    # Wrap the text based on character width and max text width
340    wrap_width = max_text_width // char_width
341    wrapped_text = textwrap.fill(recorded_error_text, width=wrap_width)
342
343    # Split the wrapped text into lines
344    lines = wrapped_text.split("\n")
345
346    # Draw each line of wrapped text
347    y = data_image.shape[0] - 300  # Starting y position
348    line_height = text_height + 10  # Spacing between lines
349    for line in lines:
350        # Put each line of text on the image
351        cv2.putText(
352            data_image,
353            line,
354            (10, y),
355            font,
356            font_scale,
357            (0, 0, 0),
358            font_thickness,
359        )
360        y += line_height  # Move down for the next line
361
362    # Log the time taken to draw the table
363    ulogger.info(f"Table Drawing Time: {time.time() - start:.2f}s")
364
365    return data_image