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
 21import tracemalloc
 22from typing import Any, BinaryIO, Callable, Dict, List, Tuple, Union
 23
 24import pydicom
 25from moviepy import VideoFileClip
 26from moviepy.video.io.ImageSequenceClip import ImageSequenceClip
 27from PIL import Image
 28from plotly.graph_objs import Figure
 29from radstract.data.dicom import (
 30    DicomTypes,
 31    convert_dicom_to_images,
 32    convert_images_to_dicom,
 33)
 34from radstract.data.nifti import NIFTI, convert_images_to_nifti_labels
 35
 36from retuve.classes.draw import Overlay
 37from retuve.classes.metrics import Metric2D, Metric3D
 38from retuve.classes.seg import SegFrameObjects
 39from retuve.hip_us.classes.dev import DevMetricsUS
 40from retuve.hip_us.classes.general import HipDatasUS, HipDataUS
 41from retuve.hip_us.draw import draw_hips_us, draw_table
 42from retuve.hip_us.handlers.bad_data import handle_bad_frames
 43from retuve.hip_us.handlers.side import set_side_info
 44from retuve.hip_us.metrics.dev import get_dev_metrics
 45from retuve.hip_us.modes.landmarks import landmarks_2_metrics_us
 46from retuve.hip_us.modes.seg import pre_process_segs_us, segs_2_landmarks_us
 47from retuve.hip_us.multiframe import (
 48    find_graf_plane,
 49    get_3d_metrics_and_visuals,
 50)
 51from retuve.hip_xray.classes import DevMetricsXRay, HipDataXray, LandmarksXRay
 52from retuve.hip_xray.draw import draw_hips_xray
 53from retuve.hip_xray.landmarks import landmarks_2_metrics_xray
 54from retuve.keyphrases.config import Config, OperationType
 55from retuve.keyphrases.enums import HipMode
 56from retuve.logs import ulogger
 57from retuve.typehints import GeneralModeFuncType
 58from retuve.custom import (
 59    custom_seg_preprocessing,
 60    get_all_custom_metrics,
 61    get_per_frame_xray,
 62)
 63
 64
 65def get_fps(no_of_frames: int, min_fps=30, min_vid_length=6) -> int:
 66    """
 67    Get the frames per second for the video clip.
 68
 69    Should be min_fps or number of fps to produce 6 min_vid_length of video.
 70
 71    :param no_of_frames: The number of frames.
 72    :param min_fps: The minimum frames per second.
 73    :param min_vid_length: The minimum video length.
 74
 75    :return: The frames per second.
 76    """
 77
 78    fps = (
 79        min_fps
 80        if no_of_frames > (min_fps * min_vid_length)
 81        else no_of_frames // min_vid_length
 82    )
 83
 84    return fps if fps > 0 else 1
 85
 86
 87def process_landmarks_xray(
 88    config: Config,
 89    landmark_results: List[LandmarksXRay],
 90    seg_results: List[SegFrameObjects],
 91) -> Tuple[List[HipDataXray], List[Image.Image]]:
 92    """
 93    Process the landmarks for the xray.
 94
 95    :param config: The configuration.
 96    :param landmark_results: The landmark results.
 97    :param seg_results: The segmentation results.
 98
 99    :return: The hip datas and the image arrays.
100    """
101    hip_datas_xray = landmarks_2_metrics_xray(landmark_results, config)
102
103    hip_datas_xray = get_per_frame_xray(hip_datas_xray, seg_results, config)
104
105    image_arrays = draw_hips_xray(hip_datas_xray, seg_results, config)
106    return hip_datas_xray, image_arrays
107
108
109def process_segs_us(
110    config: Config,
111    file: BinaryIO,
112    modes_func: Callable[
113        [BinaryIO, Union[str, Config], Dict[str, Any]],
114        List[SegFrameObjects],
115    ],
116    modes_func_kwargs_dict: Dict[str, Any],
117    called_by_2dus: bool = False,
118) -> Tuple[HipDatasUS, List[SegFrameObjects], Tuple[int, int, int]]:
119    """
120    Process the segmentation for the 3DUS.
121
122    :param config: The configuration.
123    :param file: The file.
124    :param modes_func: The mode function.
125    :param modes_func_kwargs_dict: The mode function kwargs.
126    :param called_by_2dus: Whether the calling was from the 2DUS function.
127
128    :return: The hip datas, the results, and the shape.
129    """
130
131    results: List[SegFrameObjects] = modes_func(file, config, **modes_func_kwargs_dict)
132    results, shape = pre_process_segs_us(results, config)
133
134    results = custom_seg_preprocessing(results, shape, config)
135
136    if config.test_data_passthrough:
137        pre_edited_results = copy.deepcopy(results)
138
139    landmarks, all_seg_rejection_reasons, ilium_angle_baselines = segs_2_landmarks_us(
140        results, config
141    )
142
143    if config.test_data_passthrough:
144        pre_edited_landmarks = copy.deepcopy(landmarks)
145
146    hip_datas = landmarks_2_metrics_us(landmarks, shape, config)
147    hip_datas.all_seg_rejection_reasons = all_seg_rejection_reasons
148    hip_datas.ilium_angle_baselines = ilium_angle_baselines
149
150    hip_datas = get_all_custom_metrics(
151        hip_datas, results, config, called_by_2dus=called_by_2dus
152    )
153
154    if config.test_data_passthrough:
155        hip_datas.pre_edited_results = pre_edited_results
156        hip_datas.pre_edited_landmarks = pre_edited_landmarks
157        hip_datas.pre_edited_hip_datas = copy.deepcopy(hip_datas)
158
159    return hip_datas, results, shape
160
161
162def analyse_hip_xray_2D(
163    img: Union[Image.Image, pydicom.FileDataset],
164    keyphrase: Union[str, Config],
165    modes_func: Callable[
166        [Image.Image, str, Dict[str, Any]],
167        Tuple[List[LandmarksXRay], List[SegFrameObjects]],
168    ],
169    modes_func_kwargs_dict: Dict[str, Any],
170) -> Tuple[HipDataXray, Image.Image, DevMetricsXRay]:
171    """
172    Analyze the hip for the xray.
173
174    :param img: The image.
175    :param keyphrase: The keyphrase.
176    :param modes_func: The mode function.
177    :param modes_func_kwargs_dict: The mode function kwargs.
178
179    :return: The hip, the image, and the dev metrics.
180    """
181    if isinstance(keyphrase, str):
182        config = Config.get_config(keyphrase)
183    else:
184        config = keyphrase
185
186    if isinstance(img, pydicom.FileDataset):
187        data = img
188    elif isinstance(img, Image.Image):
189        data = [img]
190    else:
191        raise ValueError(f"Invalid image type: {type(img)}. Expected Image or DICOM.")
192
193    if config.operation_type in OperationType.LANDMARK:
194        landmark_results, seg_results = modes_func(
195            data, keyphrase, **modes_func_kwargs_dict
196        )
197        hip_datas, image_arrays = process_landmarks_xray(
198            config, landmark_results, seg_results
199        )
200
201    img = image_arrays[0]
202    img = Image.fromarray(img)
203    hip = hip_datas[0]
204
205    if config.test_data_passthrough:
206        hip.seg_results = seg_results
207
208    return hip, img, DevMetricsXRay()
209
210
211def analyze_synthetic_xray(
212    dcm: pydicom.FileDataset,
213    keyphrase: Union[str, Config],
214    modes_func: Callable[
215        [pydicom.FileDataset, str, Dict[str, Any]],
216        Tuple[List[LandmarksXRay], List[SegFrameObjects]],
217    ],
218    modes_func_kwargs_dict: Dict[str, Any],
219) -> NIFTI:
220    """
221    NOTE: Experimental function.
222
223    Useful if the xray images are stacked in a single DICOM file.
224
225    Analyze the hip for the xray.
226
227    :param dcm: The DICOM file.
228    :param keyphrase: The keyphrase.
229    :param modes_func: The mode function.
230    :param modes_func_kwargs_dict: The mode function kwargs.
231
232    :return: The nifti segmentation file
233    """
234    if isinstance(keyphrase, str):
235        config = Config.get_config(keyphrase)
236    else:
237        config = keyphrase
238
239    images = convert_dicom_to_images(dcm)
240    nifti_frames = []
241
242    try:
243        if config.operation_type in OperationType.LANDMARK:
244            landmark_results, seg_results = modes_func(
245                images, keyphrase, **modes_func_kwargs_dict
246            )
247            hip_datas, image_arrays = process_landmarks_xray(
248                config, landmark_results, seg_results
249            )
250    except Exception as e:
251        if config.batch.debug == True:
252            raise e
253        ulogger.error(f"Critical Error: {e}")
254        return None
255
256    for hip, seg_frame_objs in zip(hip_datas, seg_results):
257        shape = seg_frame_objs.img.shape
258
259        overlay = Overlay((shape[0], shape[1], 3), config)
260        test = overlay.get_nifti_frame(seg_frame_objs, shape)
261        nifti_frames.append(test)
262
263    # Convert to NIfTI
264    nifti = convert_images_to_nifti_labels(nifti_frames)
265
266    return nifti
267
268
269def analyse_hip_3DUS(
270    image: Union[pydicom.FileDataset, List[Image.Image]],
271    keyphrase: Union[str, Config],
272    modes_func: Callable[
273        [pydicom.FileDataset, str, Dict[str, Any]],
274        List[SegFrameObjects],
275    ],
276    modes_func_kwargs_dict: Dict[str, Any],
277) -> Tuple[
278    HipDatasUS,
279    ImageSequenceClip,
280    Figure,
281    Union[DevMetricsXRay, DevMetricsUS],
282]:
283    """
284    Analyze a 3D Ultrasound Hip
285
286    :param dcm: The DICOM file.
287    :param keyphrase: The keyphrase.
288    :param modes_func: The mode function.
289    :param modes_func_kwargs_dict: The mode function kwargs.
290
291    :return: The hip datas, the video clip, the 3D visual, and the dev metrics.
292    """
293    start = time.time()
294
295    config = Config.get_config(keyphrase)
296    hip_datas = HipDatasUS()
297
298    file_id = modes_func_kwargs_dict.get("file_id")
299    if file_id:
300        del modes_func_kwargs_dict["file_id"]
301
302    # if a set of images, convert to a DICOM file
303    if isinstance(image, list) and all(isinstance(img, Image.Image) for img in image):
304        image = convert_images_to_dicom(image)
305
306    try:
307        if config.operation_type == OperationType.SEG:
308            hip_datas, results, shape = process_segs_us(
309                config, image, modes_func, modes_func_kwargs_dict
310            )
311        elif config.operation_type == OperationType.LANDMARK:
312            raise NotImplementedError(
313                "This is not yet supported. Please use the seg operation type."
314            )
315    except Exception as e:
316        if config.batch.debug == True:
317            raise e
318        ulogger.error(f"Critical Error: {e}")
319        return None, None, None, None
320
321    hip_datas = handle_bad_frames(hip_datas, config)
322
323    if not any(hip.metrics for hip in hip_datas):
324        ulogger.error(f"No metrics were found in image.")
325
326    hip_datas.file_id = file_id
327    hip_datas = find_graf_plane(hip_datas, results, config=config)
328
329    hip_datas, results = set_side_info(hip_datas, results, config)
330
331    (
332        hip_datas,
333        visual_3d,
334        fem_sph,
335        illium_mesh,
336        apex_points,
337        femoral_sphere,
338        avg_normals_data,
339        normals_data,
340    ) = get_3d_metrics_and_visuals(hip_datas, results, config)
341
342    image_arrays, nifti = draw_hips_us(hip_datas, results, fem_sph, config)
343
344    if config.seg_export:
345        hip_datas.nifti = nifti
346
347    hip_datas = get_dev_metrics(hip_datas, results, config)
348
349    # data_image = draw_table(shape, hip_datas)
350    # image_arrays.append(data_image)
351
352    ulogger.info(f"Total 3DUS time: {time.time() - start:.2f}s")
353
354    fps = get_fps(
355        len(image_arrays),
356        config.visuals.min_vid_fps,
357        config.visuals.min_vid_length,
358    )
359
360    video_clip = ImageSequenceClip(
361        image_arrays,
362        fps=fps,
363    )
364
365    if config.test_data_passthrough:
366        hip_datas.illium_mesh = illium_mesh
367        hip_datas.fem_sph = fem_sph
368        hip_datas.results = results
369        hip_datas.apex_points = apex_points
370        hip_datas.femoral_sphere = femoral_sphere
371        hip_datas.avg_normals_data = avg_normals_data
372        hip_datas.normals_data = normals_data
373
374    if hip_datas.custom_metrics is not None:
375        hip_datas.metrics += hip_datas.custom_metrics
376
377    return (
378        hip_datas,
379        video_clip,
380        visual_3d,
381        hip_datas.dev_metrics,
382    )
383
384
385def analyse_hip_2DUS(
386    img: Union[Image.Image, pydicom.FileDataset],
387    keyphrase: Union[str, Config],
388    modes_func: Callable[
389        [Image.Image, str, Dict[str, Any]],
390        List[SegFrameObjects],
391    ],
392    modes_func_kwargs_dict: Dict[str, Any],
393    return_seg_info: bool = False,
394) -> Tuple[HipDataUS, Image.Image, DevMetricsUS]:
395    """
396    Analyze a 2D Ultrasound Hip
397
398    :param img: The image.
399    :param keyphrase: The keyphrase.
400    :param modes_func: The mode function.
401    :param modes_func_kwargs_dict: The mode function kwargs.
402
403    :return: The hip, the image, and the dev metrics.
404    """
405    config = Config.get_config(keyphrase)
406
407    if isinstance(img, pydicom.FileDataset):
408        data = img
409    elif isinstance(img, Image.Image):
410        data = [img]
411
412    try:
413        if config.operation_type in OperationType.SEG:
414            hip_datas, results, _ = process_segs_us(
415                config,
416                data,
417                modes_func,
418                modes_func_kwargs_dict,
419                called_by_2dus=True,
420            )
421    except Exception as e:
422        if config.batch.debug == True:
423            raise e
424        ulogger.error(f"Critical Error: {e}")
425        return None, None, None
426
427    image_arrays, _ = draw_hips_us(hip_datas, results, None, config)
428
429    hip_datas = get_dev_metrics(hip_datas, results, config)
430
431    image = image_arrays[0]
432    hip = hip_datas[0]
433
434    image = Image.fromarray(image)
435
436    if return_seg_info:
437        hip.seg_info = results
438
439    if hip_datas.custom_metrics is not None:
440        hip.metrics += hip_datas.custom_metrics
441
442    return hip, image, hip_datas.dev_metrics
443
444
445def analyse_hip_2DUS_sweep(
446    image: Union[pydicom.FileDataset, List[Image.Image]],
447    keyphrase: Union[str, Config],
448    modes_func: Callable[
449        [pydicom.FileDataset, str, Dict[str, Any]],
450        List[SegFrameObjects],
451    ],
452    modes_func_kwargs_dict: Dict[str, Any],
453) -> Tuple[HipDataUS, Image.Image, DevMetricsUS, ImageSequenceClip]:
454
455    config = Config.get_config(keyphrase)
456    config.batch.hip_mode = HipMode.US2DSW
457    hip_datas = HipDatasUS()
458
459    # Convert list of images to DICOM
460    if isinstance(image, list) and all(isinstance(img, Image.Image) for img in image):
461        image = convert_images_to_dicom(image)
462
463    try:
464        if config.operation_type == OperationType.SEG:
465            hip_datas, results, shape = process_segs_us(
466                config,
467                image,
468                modes_func,
469                modes_func_kwargs_dict,
470                called_by_2dus=True,
471            )
472        else:
473            raise NotImplementedError("Only SEG operation type supported.")
474    except Exception as e:
475        if config.batch.debug == True:
476            raise e
477        ulogger.error(f"Critical Error: {e}")
478        return None, None, None, None
479
480    hip_datas = handle_bad_frames(hip_datas, config)
481    hip_datas = find_graf_plane(hip_datas, results, config)
482
483    graf_hip = hip_datas.grafs_hip
484    graf_frame = hip_datas.graf_frame
485    graf_hip.graf_frame = graf_frame
486    graf_hip.recorded_error = hip_datas.recorded_error
487
488    image_arrays, _ = draw_hips_us(hip_datas, results, None, config)
489
490    hip_datas = get_dev_metrics(hip_datas, results, config)
491
492    if graf_frame is not None:
493        graf_image = Image.fromarray(image_arrays[graf_frame])
494    else:
495        graf_image = Image.fromarray(image_arrays[len(image_arrays) // 2])
496        marked_pairs = [
497            (hip, Image.fromarray(image), conf)
498            for hip, image, conf in zip(hip_datas, image_arrays, hip_datas.graf_confs)
499            if hip.marked
500        ]
501
502        try:
503            if marked_pairs:
504                graf_image = marked_pairs[len(marked_pairs) // 2][1]
505
506                graf_image = min(
507                    marked_pairs,
508                    key=lambda pair: (
509                        abs(pair[0].landmarks.left[1] - pair[0].landmarks.apex[1]),
510                        -abs(pair[0].landmarks.apex[0] - pair[0].landmarks.left[0]),
511                    ),
512                )[1]
513
514                graf_image = max(
515                    marked_pairs,
516                    key=lambda pair: (pair[2]),
517                )[1]
518        except AttributeError:
519            pass
520
521    video_clip = ImageSequenceClip(
522        image_arrays,
523        fps=get_fps(
524            len(image_arrays),
525            config.visuals.min_vid_fps,
526            config.visuals.min_vid_length,
527        ),
528    )
529
530    if hip_datas.custom_metrics is not None and graf_hip.metrics:
531        graf_hip.metrics += hip_datas.custom_metrics
532
533    return graf_hip, graf_image, hip_datas.dev_metrics, video_clip
534
535
536class RetuveResult:
537    """
538    The standardised result of the Retuve pipeline.
539
540    :attr hip_datas: The hip datas.
541    :attr hip: The hip.
542    :attr image: The saved image, if any.
543    :attr metrics: The metrics.
544    :attr video_clip: The video clip, if any.
545    :attr visual_3d: The 3D visual, if any.
546    """
547
548    def __init__(
549        self,
550        metrics: Union[List[Metric2D], List[Metric3D]],
551        hip_datas: Union[HipDatasUS, List[HipDataXray]] = None,
552        hip: Union[HipDataXray, HipDataUS] = None,
553        image: Image.Image = None,
554        video_clip: ImageSequenceClip = None,
555        visual_3d: Figure = None,
556    ):
557        self.hip_datas = hip_datas
558        self.hip = hip
559        self.image = image
560        self.metrics = metrics
561        self.video_clip = video_clip
562        self.visual_3d = visual_3d
563
564
565def retuve_run(
566    hip_mode: HipMode,
567    config: Config,
568    modes_func: GeneralModeFuncType,
569    modes_func_kwargs_dict: Dict[str, Any],
570    file: str,
571) -> RetuveResult:
572    org_file_name = file
573    # 0 or 1 because we assume nothing means no extention dicoms
574    always_dcm = (
575        len(config.batch.input_types) == 1 and ".dcm" in config.batch.input_types
576    ) or config.batch.input_types == [""]
577
578    is_dicom = always_dcm or (
579        file.endswith(".dcm") and ".dcm" in config.batch.input_types
580    )
581
582    # Helper to load an image (DICOM or regular)
583    def load_image(path: str):
584        if is_dicom:
585            with pydicom.dcmread(path) as ds:
586                image = convert_dicom_to_images(ds, dicom_type=DicomTypes.SINGLE)[
587                    0
588                ].convert("RGB")
589            return image
590        else:
591            return Image.open(path).convert("RGB")
592
593    if hip_mode == HipMode.XRAY:
594        if is_dicom:
595            with pydicom.dcmread(file) as ds:
596                file_obj = ds
597                hip, image, dev_metrics = analyse_hip_xray_2D(
598                    file_obj, config, modes_func, modes_func_kwargs_dict
599                )
600        else:
601            img = Image.open(file).convert("RGB")
602            hip, image, dev_metrics = analyse_hip_xray_2D(
603                img, config, modes_func, modes_func_kwargs_dict
604            )
605
606            dump = hip.json_dump(config, dev_metrics)
607            dump["landmarks"] = dict(hip.landmarks.items())
608
609        return RetuveResult(dump, image=image, hip=hip)
610
611    elif hip_mode == HipMode.US2D:
612        img = load_image(file)
613        hip, image, dev_metrics = analyse_hip_2DUS(
614            img, config, modes_func, modes_func_kwargs_dict
615        )
616        return RetuveResult(hip.json_dump(config, dev_metrics), hip=hip, image=image)
617
618    elif hip_mode == HipMode.US2DSW:
619        if is_dicom:
620            with pydicom.dcmread(file) as ds:
621                file_obj = ds
622                hip, image, dev_metrics, video_clip = analyse_hip_2DUS_sweep(
623                    file_obj, config, modes_func, modes_func_kwargs_dict
624                )
625        else:
626            if ".mp4" in file:
627                file = video_to_pillow_images(file)
628
629            hip, image, dev_metrics, video_clip = analyse_hip_2DUS_sweep(
630                file, config, modes_func, modes_func_kwargs_dict
631            )
632
633        json_dump = hip.json_dump(config, dev_metrics) if hip else None
634        return RetuveResult(
635            json_dump,
636            hip=hip,
637            image=image,
638            video_clip=video_clip,
639        )
640
641    elif hip_mode == HipMode.US3D:
642        modes_func_kwargs_dict["file_id"] = org_file_name.split("/")[-1]
643
644        if is_dicom:
645            with pydicom.dcmread(file) as ds:
646                hip_datas, video_clip, visual_3d, dev_metrics = analyse_hip_3DUS(
647                    ds, config, modes_func, modes_func_kwargs_dict
648                )
649        else:
650            hip_datas, video_clip, visual_3d, dev_metrics = analyse_hip_3DUS(
651                file, config, modes_func, modes_func_kwargs_dict
652            )
653
654        if hip_datas:
655            return RetuveResult(
656                hip_datas.json_dump(config),
657                hip_datas=hip_datas,
658                video_clip=video_clip,
659                visual_3d=visual_3d,
660            )
661        else:
662            return RetuveResult({})
663
664    else:
665        raise ValueError(f"Invalid hip_mode: {hip_mode}")
666
667
668def video_to_pillow_images(video_path):
669    """
670    Opens a video file and converts each frame into a Pillow Image object.
671
672    Args:
673        video_path (str): The path to the video file.
674
675    Returns:
676        list: A list of Pillow Image objects, one for each frame of the video.
677    """
678    try:
679        clip = VideoFileClip(video_path)
680
681        image_list = []
682        for frame_np_array in clip.iter_frames():
683            # Convert the NumPy array (frame) to a Pillow Image
684            image = Image.fromarray(frame_np_array)
685            image_list.append(image)
686
687        clip.close()  # Close the clip to release resources
688        return image_list
689
690    except Exception as e:
691        print(f"An error occurred: {e}")
692        return []