retuve.funcs

  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"""
 16Contains the high-level functions that are used to run the Retuve pipeline.
 17"""
 18
 19import copy
 20import time
 21from typing import Any, BinaryIO, Callable, Dict, List, Tuple, Union
 22
 23import pydicom
 24from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
 25from PIL import Image
 26from plotly.graph_objs import Figure
 27from radstract.data.dicom import convert_dicom_to_images
 28from radstract.data.nifti import NIFTI, convert_images_to_nifti_labels
 29
 30from retuve.classes.draw import Overlay
 31from retuve.classes.metrics import Metric2D, Metric3D
 32from retuve.classes.seg import SegFrameObjects
 33from retuve.hip_us.classes.dev import DevMetricsUS
 34from retuve.hip_us.classes.general import HipDatasUS, HipDataUS
 35from retuve.hip_us.draw import draw_hips_us, draw_table
 36from retuve.hip_us.handlers.bad_data import handle_bad_frames
 37from retuve.hip_us.handlers.side import set_side
 38from retuve.hip_us.metrics.dev import get_dev_metrics
 39from retuve.hip_us.modes.landmarks import landmarks_2_metrics_us
 40from retuve.hip_us.modes.seg import pre_process_segs_us, segs_2_landmarks_us
 41from retuve.hip_us.multiframe import (find_graf_plane,
 42                                      get_3d_metrics_and_visuals)
 43from retuve.hip_xray.classes import DevMetricsXRay, HipDataXray, LandmarksXRay
 44from retuve.hip_xray.draw import draw_hips_xray
 45from retuve.hip_xray.landmarks import landmarks_2_metrics_xray
 46from retuve.keyphrases.config import Config, OperationType
 47from retuve.keyphrases.enums import HipMode
 48from retuve.logs import ulogger
 49from retuve.typehints import GeneralModeFuncType
 50
 51
 52def get_fps(no_of_frames: int, min_fps=30, min_vid_length=6) -> int:
 53    """
 54    Get the frames per second for the video clip.
 55
 56    Should be min_fps or number of fps to produce 6 min_vid_length of video.
 57
 58    :param no_of_frames: The number of frames.
 59    :param min_fps: The minimum frames per second.
 60    :param min_vid_length: The minimum video length.
 61
 62    :return: The frames per second.
 63    """
 64
 65    fps = (
 66        min_fps
 67        if no_of_frames > (min_fps * min_vid_length)
 68        else no_of_frames // min_vid_length
 69    )
 70
 71    return fps if fps > 0 else 1
 72
 73
 74def process_landmarks_xray(
 75    config: Config,
 76    landmark_results: List[LandmarksXRay],
 77    seg_results: List[SegFrameObjects],
 78) -> Tuple[List[HipDataXray], List[Image.Image]]:
 79    """
 80    Process the landmarks for the xray.
 81
 82    :param config: The configuration.
 83    :param landmark_results: The landmark results.
 84    :param seg_results: The segmentation results.
 85
 86    :return: The hip datas and the image arrays.
 87    """
 88    hip_datas_xray = landmarks_2_metrics_xray(landmark_results, config)
 89    image_arrays = draw_hips_xray(hip_datas_xray, seg_results, config)
 90    return hip_datas_xray, image_arrays
 91
 92
 93def process_segs_us(
 94    config: Config,
 95    file: BinaryIO,
 96    modes_func: Callable[
 97        [BinaryIO, Union[str, Config], Dict[str, Any]],
 98        List[SegFrameObjects],
 99    ],
100    modes_func_kwargs_dict: Dict[str, Any],
101) -> Tuple[HipDatasUS, List[SegFrameObjects], Tuple[int, int, int]]:
102    """
103    Process the segmentation for the 3DUS.
104
105    :param config: The configuration.
106    :param file: The file.
107    :param modes_func: The mode function.
108    :param modes_func_kwargs_dict: The mode function kwargs.
109
110    :return: The hip datas, the results, and the shape.
111    """
112
113    results: List[SegFrameObjects] = modes_func(
114        file, config, **modes_func_kwargs_dict
115    )
116    results, shape = pre_process_segs_us(results, config)
117    pre_edited_results = copy.deepcopy(results)
118    landmarks, all_seg_rejection_reasons = segs_2_landmarks_us(results, config)
119    pre_edited_landmarks = copy.deepcopy(landmarks)
120    hip_datas = landmarks_2_metrics_us(landmarks, shape, config)
121    hip_datas.all_seg_rejection_reasons = all_seg_rejection_reasons
122
123    if config.test_data_passthrough:
124        hip_datas.pre_edited_results = pre_edited_results
125        hip_datas.pre_edited_landmarks = pre_edited_landmarks
126        hip_datas.pre_edited_hip_datas = copy.deepcopy(hip_datas)
127
128    return hip_datas, results, shape
129
130
131def analyse_hip_xray_2D(
132    img: Image.Image,
133    keyphrase: Union[str, Config],
134    modes_func: Callable[
135        [Image.Image, str, Dict[str, Any]],
136        Tuple[List[LandmarksXRay], List[SegFrameObjects]],
137    ],
138    modes_func_kwargs_dict: Dict[str, Any],
139) -> Tuple[HipDataXray, Image.Image, DevMetricsXRay]:
140    """
141    Analyze the hip for the xray.
142
143    :param img: The image.
144    :param keyphrase: The keyphrase.
145    :param modes_func: The mode function.
146    :param modes_func_kwargs_dict: The mode function kwargs.
147
148    :return: The hip, the image, and the dev metrics.
149    """
150    if isinstance(keyphrase, str):
151        config = Config.get_config(keyphrase)
152    else:
153        config = keyphrase
154
155    if config.operation_type in OperationType.LANDMARK:
156        landmark_results, seg_results = modes_func(
157            [img], keyphrase, **modes_func_kwargs_dict
158        )
159        hip_datas, image_arrays = process_landmarks_xray(
160            config, landmark_results, seg_results
161        )
162
163    img = image_arrays[0]
164    img = Image.fromarray(img)
165    hip = hip_datas[0]
166
167    if config.test_data_passthrough:
168        hip.seg_results = seg_results
169
170    return hip, img, DevMetricsXRay()
171
172
173def analyze_synthetic_xray(
174    dcm: pydicom.FileDataset,
175    keyphrase: Union[str, Config],
176    modes_func: Callable[
177        [pydicom.FileDataset, str, Dict[str, Any]],
178        Tuple[List[LandmarksXRay], List[SegFrameObjects]],
179    ],
180    modes_func_kwargs_dict: Dict[str, Any],
181) -> NIFTI:
182    """
183    NOTE: Experimental function.
184
185    Useful if the xray images are stacked in a single DICOM file.
186
187    Analyze the hip for the xray.
188
189    :param dcm: The DICOM file.
190    :param keyphrase: The keyphrase.
191    :param modes_func: The mode function.
192    :param modes_func_kwargs_dict: The mode function kwargs.
193
194    :return: The nifti segmentation file
195    """
196    if isinstance(keyphrase, str):
197        config = Config.get_config(keyphrase)
198    else:
199        config = keyphrase
200
201    images = convert_dicom_to_images(dcm)
202    nifti_frames = []
203
204    try:
205        if config.operation_type in OperationType.LANDMARK:
206            landmark_results, seg_results = modes_func(
207                images, keyphrase, **modes_func_kwargs_dict
208            )
209            hip_datas, image_arrays = process_landmarks_xray(
210                config, landmark_results, seg_results
211            )
212    except Exception as e:
213        ulogger.error(f"Critical Error: {e}")
214        return None
215
216    for hip, seg_frame_objs in zip(hip_datas, seg_results):
217        shape = seg_frame_objs.img.shape
218
219        overlay = Overlay((shape[0], shape[1], 3), config)
220        test = overlay.get_nifti_frame(seg_frame_objs, shape)
221        nifti_frames.append(test)
222
223    # Convert to NIfTI
224    nifti = convert_images_to_nifti_labels(nifti_frames)
225
226    return nifti
227
228
229def analyse_hip_3DUS(
230    dcm: pydicom.FileDataset,
231    keyphrase: Union[str, Config],
232    modes_func: Callable[
233        [pydicom.FileDataset, str, Dict[str, Any]],
234        List[SegFrameObjects],
235    ],
236    modes_func_kwargs_dict: Dict[str, Any],
237) -> Tuple[
238    HipDatasUS,
239    ImageSequenceClip,
240    Figure,
241    Union[DevMetricsXRay, DevMetricsUS],
242]:
243    """
244    Analyze a 3D Ultrasound Hip
245
246    :param dcm: The DICOM file.
247    :param keyphrase: The keyphrase.
248    :param modes_func: The mode function.
249    :param modes_func_kwargs_dict: The mode function kwargs.
250
251    :return: The hip datas, the video clip, the 3D visual, and the dev metrics.
252    """
253    start = time.time()
254
255    config = Config.get_config(keyphrase)
256    hip_datas = HipDatasUS()
257
258    file_id = modes_func_kwargs_dict.get("file_id")
259    if file_id:
260        del modes_func_kwargs_dict["file_id"]
261
262    try:
263        if config.operation_type == OperationType.SEG:
264            hip_datas, results, shape = process_segs_us(
265                config, dcm, modes_func, modes_func_kwargs_dict
266            )
267        elif config.operation_type == OperationType.LANDMARK:
268            raise NotImplementedError(
269                "This is not yet supported. Please use the seg operation type."
270            )
271    except Exception as e:
272        ulogger.error(f"Critical Error: {e}")
273        return None, None, None, None
274
275    hip_datas = handle_bad_frames(hip_datas, config)
276
277    if not any(hip.metrics for hip in hip_datas):
278        dcm_patient = dcm.get("PatientID", "Unknown")
279        ulogger.error(f"No metrics were found in the DICOM {dcm_patient}.")
280
281    hip_datas.file_id = file_id
282    hip_datas = find_graf_plane(hip_datas, results, config=config)
283
284    hip_datas, results = set_side(
285        hip_datas, results, config.hip.allow_flipping
286    )
287
288    (
289        hip_datas,
290        visual_3d,
291        fem_sph,
292        illium_mesh,
293        apex_points,
294        femoral_sphere,
295        avg_normals_data,
296        normals_data,
297    ) = get_3d_metrics_and_visuals(hip_datas, results, config)
298
299    image_arrays, nifti = draw_hips_us(hip_datas, results, fem_sph, config)
300
301    if config.seg_export:
302        hip_datas.nifti = nifti
303
304    hip_datas = get_dev_metrics(hip_datas, results, config)
305
306    data_image = draw_table(shape, hip_datas)
307    image_arrays.append(data_image)
308
309    ulogger.info(f"Total 3DUS time: {time.time() - start:.2f}s")
310
311    fps = get_fps(
312        len(image_arrays),
313        config.visuals.min_vid_fps,
314        config.visuals.min_vid_length,
315    )
316
317    video_clip = ImageSequenceClip(
318        image_arrays,
319        fps=fps,
320    )
321
322    if config.test_data_passthrough:
323        hip_datas.illium_mesh = illium_mesh
324        hip_datas.fem_sph = fem_sph
325        hip_datas.results = results
326        hip_datas.apex_points = apex_points
327        hip_datas.femoral_sphere = femoral_sphere
328        hip_datas.avg_normals_data = avg_normals_data
329        hip_datas.normals_data = normals_data
330
331    return (
332        hip_datas,
333        video_clip,
334        visual_3d,
335        hip_datas.dev_metrics,
336    )
337
338
339def analyse_hip_2DUS(
340    img: Image.Image,
341    keyphrase: Union[str, Config],
342    modes_func: Callable[
343        [Image.Image, str, Dict[str, Any]],
344        List[SegFrameObjects],
345    ],
346    modes_func_kwargs_dict: Dict[str, Any],
347    return_seg_info: bool = False,
348) -> Tuple[HipDataUS, Image.Image, DevMetricsUS]:
349    """
350    Analyze a 2D Ultrasound Hip
351
352    :param img: The image.
353    :param keyphrase: The keyphrase.
354    :param modes_func: The mode function.
355    :param modes_func_kwargs_dict: The mode function kwargs.
356
357    :return: The hip, the image, and the dev metrics.
358    """
359    config = Config.get_config(keyphrase)
360
361    try:
362        if config.operation_type in OperationType.SEG:
363            hip_datas, results, _ = process_segs_us(
364                config, [img], modes_func, modes_func_kwargs_dict
365            )
366    except Exception as e:
367        ulogger.error(f"Critical Error: {e}")
368        return None, None, None
369
370    image_arrays, _ = draw_hips_us(hip_datas, results, None, config)
371
372    hip_datas = get_dev_metrics(hip_datas, results, config)
373
374    image = image_arrays[0]
375    hip = hip_datas[0]
376
377    image = Image.fromarray(image)
378
379    if return_seg_info:
380        hip.seg_info = results
381
382    return hip, image, hip_datas.dev_metrics
383
384
385def analyse_hip_2DUS_sweep(
386    dcm: pydicom.FileDataset,
387    keyphrase: Union[str, Config],
388    modes_func: Callable[
389        [pydicom.FileDataset, str, Dict[str, Any]],
390        List[SegFrameObjects],
391    ],
392    modes_func_kwargs_dict: Dict[str, Any],
393) -> Tuple[
394    HipDataUS,
395    Image.Image,
396    DevMetricsUS,
397    ImageSequenceClip,
398]:
399    """
400    Analyze a 2D Sweep Ultrasound Hip
401
402    :param dcm: The DICOM file.
403    :param keyphrase: The keyphrase.
404    :param modes_func: The mode function.
405    :param modes_func_kwargs_dict: The mode function kwargs.
406
407    :return: The hip, the image, the dev metrics, and the video clip.
408    """
409
410    config = Config.get_config(keyphrase)
411    hip_datas = HipDatasUS()
412
413    try:
414        if config.operation_type == OperationType.SEG:
415            hip_datas, results, shape = process_segs_us(
416                config, dcm, modes_func, modes_func_kwargs_dict
417            )
418        elif config.operation_type == OperationType.LANDMARK:
419            raise NotImplementedError(
420                "This is not yet supported. Please use the seg operation type."
421            )
422    except Exception as e:
423        ulogger.error(f"Critical Error: {e}")
424        return None, None, None, None
425
426    hip_datas = handle_bad_frames(hip_datas, config)
427    hip_datas = find_graf_plane(hip_datas, results, config)
428
429    graf_hip = hip_datas.grafs_hip
430    graf_frame = hip_datas.graf_frame
431
432    image_arrays, _ = draw_hips_us(hip_datas, results, None, config)
433
434    hip_datas = get_dev_metrics(hip_datas, results, config)
435
436    if graf_frame is not None:
437        graf_image = image_arrays[graf_frame]
438        graf_image = Image.fromarray(graf_image)
439
440        image_arrays = (
441            [image_arrays[graf_frame]] * int(len(image_arrays) * 0.1)
442            + image_arrays
443            + [image_arrays[graf_frame]] * int(len(image_arrays) * 0.1)
444        )
445    else:
446        graf_image = None
447
448    video_clip = ImageSequenceClip(
449        image_arrays,
450        fps=get_fps(
451            len(image_arrays),
452            config.visuals.min_vid_fps,
453            config.visuals.min_vid_length,
454        ),
455    )
456
457    return graf_hip, graf_image, hip_datas.dev_metrics, video_clip
458
459
460class RetuveResult:
461    """
462    The standardised result of the Retuve pipeline.
463
464    :attr hip_datas: The hip datas.
465    :attr hip: The hip.
466    :attr image: The saved image, if any.
467    :attr metrics: The metrics.
468    :attr video_clip: The video clip, if any.
469    :attr visual_3d: The 3D visual, if any.
470    """
471
472    def __init__(
473        self,
474        metrics: Union[List[Metric2D], List[Metric3D]],
475        hip_datas: Union[HipDatasUS, List[HipDataXray]] = None,
476        hip: Union[HipDataXray, HipDataUS] = None,
477        image: Image.Image = None,
478        video_clip: ImageSequenceClip = None,
479        visual_3d: Figure = None,
480    ):
481        self.hip_datas = hip_datas
482        self.hip = hip
483        self.image = image
484        self.metrics = metrics
485        self.video_clip = video_clip
486        self.visual_3d = visual_3d
487
488
489def retuve_run(
490    hip_mode: HipMode,
491    config: Config,
492    modes_func: GeneralModeFuncType,
493    modes_func_kwargs_dict: Dict[str, Any],
494    file: str,
495) -> RetuveResult:
496    """
497    Run the Retuve pipeline with standardised inputs and outputs.
498
499    :param hip_mode: The hip mode.
500    :param config: The configuration.
501    :param modes_func: The mode function.
502    :param modes_func_kwargs_dict: The mode function kwargs.
503    :param file: The file.
504
505    :return: The Retuve result standardised output.
506    """
507    always_dcm = (
508        len(config.batch.input_types) == 1
509        and ".dcm" in config.batch.input_types
510    )
511
512    if always_dcm or (
513        file.endswith(".dcm") and ".dcm" in config.batch.input_types
514    ):
515        file = pydicom.dcmread(file)
516
517    if hip_mode == HipMode.XRAY:
518        file = Image.open(file)
519        hip, image, dev_metrics = analyse_hip_xray_2D(
520            file, config, modes_func, modes_func_kwargs_dict
521        )
522        return RetuveResult(
523            hip.json_dump(config, dev_metrics), image=image, hip=hip
524        )
525    elif hip_mode == HipMode.US3D:
526        hip_datas, video_clip, visual_3d, dev_metrics = analyse_hip_3DUS(
527            file, config, modes_func, modes_func_kwargs_dict
528        )
529        return RetuveResult(
530            hip_datas.json_dump(config),
531            hip_datas=hip_datas,
532            video_clip=video_clip,
533            visual_3d=visual_3d,
534        )
535    elif hip_mode == HipMode.US2D:
536        file = Image.open(file).convert("RGB")
537        hip, image, dev_metrics = analyse_hip_2DUS(
538            file, config, modes_func, modes_func_kwargs_dict
539        )
540        return RetuveResult(
541            hip.json_dump(config, dev_metrics), hip=hip, image=image
542        )
543    elif hip_mode == HipMode.US2DSW:
544        hip, image, dev_metrics, video_clip = analyse_hip_2DUS_sweep(
545            file, config, modes_func, modes_func_kwargs_dict
546        )
547        json_dump = None
548        if hip:
549            json_dump = hip.json_dump(config, dev_metrics)
550
551        return RetuveResult(
552            json_dump,
553            hip=hip,
554            image=image,
555            video_clip=video_clip,
556        )
557    else:
558        raise ValueError(f"Invalid hip_mode. {hip_mode}")