retuve.hip_us.modes.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"""
 16All code relating to going from segmentation results to landmarks is in this file.
 17"""
 18
 19import copy
 20import time
 21from typing import List
 22
 23import cv2
 24import numpy as np
 25from skimage.morphology import skeletonize
 26
 27from retuve.classes.seg import (
 28    MidLine,
 29    NDArrayImg_NxNx3_AllWhite,
 30    SegFrameObjects,
 31    SegObject,
 32)
 33from retuve.hip_us.classes.enums import HipLabelsUS
 34from retuve.hip_us.classes.general import LandmarksUS
 35from retuve.hip_us.handlers.segs import remove_bad_objs
 36from retuve.hip_us.metrics.alpha import find_alpha_landmarks
 37from retuve.hip_us.metrics.coverage import find_cov_landmarks
 38from retuve.keyphrases.config import Config
 39from retuve.keyphrases.enums import MidLineMove
 40from retuve.logs import log_timings
 41from retuve.utils import find_midline_extremes
 42
 43
 44def get_midlines(
 45    mask: NDArrayImg_NxNx3_AllWhite, config: Config, color: int = 255
 46) -> MidLine:
 47    """
 48    Get the midline of the illium mask. This is used to '
 49    find the landmarks for the alpha angle.
 50
 51    :param mask: The mask of the illium.
 52    :param config: The configuration object.
 53    :param color: The color of the midline.
 54
 55    :return: The midline of the illium.
 56    """
 57
 58    mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
 59    # Takes bool as input and returns bool as output
 60    midline = skeletonize(mask > 0)
 61    midline = midline.astype(np.uint8) * color
 62
 63    midline_moved = midline
 64
 65    # move the skeleton up and right
 66    if config.hip.midline_move_method == MidLineMove.BASIC:
 67        height, width = midline_moved.shape
 68        midline_moved = np.roll(midline_moved, -1 * int(height / 51), axis=0)
 69        midline_moved = np.roll(midline_moved, int(width / 86), axis=1)
 70
 71    for midl in [midline_moved, midline]:
 72        left_most, right_most = find_midline_extremes(
 73            np.column_stack(np.where(midl == color))
 74        )
 75        if left_most is not None and right_most is not None:
 76            width = right_most[1] - left_most[1]
 77
 78            new_left_bound = max(int(left_most[1] + width * 0.01), 4)
 79            new_right_bound = max(int(right_most[1] - width * 0.01), 4)
 80
 81            # Apply the boundary conditions to remove edges
 82            midl[:, :new_left_bound] = 0
 83            midl[:, new_right_bound:] = 0
 84
 85    midline = np.column_stack(np.where(midline == color))
 86    midline_moved = np.column_stack(np.where(midline_moved == color))
 87
 88    return midline, midline_moved
 89
 90
 91def segs_2_landmarks_us(
 92    results: List[SegFrameObjects], config: Config
 93) -> List[LandmarksUS]:
 94    """
 95    Converts a list of segmentation results to a list of landmarks.
 96
 97    :param results: A list of segmentation results.
 98    :param config: The configuration object.
 99
100    :return: A list of landmarks.
101    """
102    hip_landmarks = []
103    timings = []
104    hip_objs_list = []
105    fem_head_ilium_wrong_way_round = 0
106    all_rejection_reasons = []
107
108    for seg_frame_objs in results:
109
110        if all(seg_obj.empty for seg_obj in seg_frame_objs):
111            hip_objs_list.append(None)
112            all_rejection_reasons.append([])
113            continue
114
115        hip_objs = {
116            HipLabelsUS.IlliumAndAcetabulum: [],
117            HipLabelsUS.FemoralHead: [],
118            HipLabelsUS.OsIchium: [],
119        }
120
121        for seg_obj in seg_frame_objs:
122            hip_objs[seg_obj.cls].append(seg_obj)
123
124        hip_objs, wrong_way_round, rejection_reasons = remove_bad_objs(
125            hip_objs, seg_frame_objs.img
126        )
127        if wrong_way_round:
128            fem_head_ilium_wrong_way_round += 1
129
130        hip_objs_list.append(hip_objs)
131        all_rejection_reasons.append(rejection_reasons)
132
133    if fem_head_ilium_wrong_way_round > 5:
134        # Flip images
135        for seg_frame_objs in results:
136            seg_frame_objs.img = cv2.flip(seg_frame_objs.img, 1)
137
138        for hip_objs in hip_objs_list:
139            if hip_objs is None:
140                continue
141
142            for seg_obj in hip_objs.values():
143                seg_obj.flip_horizontally(results[0].img.shape[1])
144
145            ilium = hip_objs.get(HipLabelsUS.IlliumAndAcetabulum, None)
146
147            if (
148                ilium is not None
149                and ilium.box is not None
150                and ilium.box[0] > results[0].img.shape[1] / 2
151            ):
152                hip_objs[HipLabelsUS.IlliumAndAcetabulum] = SegObject(
153                    empty=True
154                )
155
156    for seg_frame_objs, hip_objs in zip(results, hip_objs_list):
157        start = time.time()
158
159        landmarks = LandmarksUS()
160        if hip_objs is None:
161            hip_landmarks.append(landmarks)
162            continue
163
164        # set the seg_frame_objs to only have the good seg_objs
165        seg_frame_objs.seg_objects = [
166            obj for obj in list(hip_objs.values()) if obj is not None
167        ]
168
169        illium = hip_objs.get(HipLabelsUS.IlliumAndAcetabulum, None)
170        femoral = hip_objs.get(HipLabelsUS.FemoralHead, None)
171
172        landmarks = find_alpha_landmarks(illium, landmarks, config)
173        landmarks = find_cov_landmarks(femoral, landmarks, config)
174
175        hip_landmarks.append(landmarks)
176        timings.append(time.time() - start)
177
178    log_timings(timings, title="Seg->Landmarking Speed:")
179    return hip_landmarks, all_rejection_reasons
180
181
182def pre_process_segs_us(
183    results: List[SegFrameObjects], config: Config
184) -> tuple[List[SegFrameObjects], tuple]:
185    """
186    Pre-processes the segmentation results for the hip US module.
187
188    :param results: A list of segmentation results.
189    :param config: The configuration object.
190
191    :return: A tuple containing the pre-processed
192             segmentation results and the shape of the image.
193    """
194    timings = []
195    for frame_seg_objs in results:
196        start = time.time()
197        for seg_object in frame_seg_objs:
198            if (
199                seg_object.empty
200                or seg_object.cls != HipLabelsUS.IlliumAndAcetabulum
201            ):
202                continue
203
204            seg_object.midline, seg_object.midline_moved = get_midlines(
205                seg_object.mask, config
206            )
207
208        timings.append(time.time() - start)
209
210    log_timings(timings, title="Skeletonization Speed:")
211
212    img = results[0].img
213    shape = img.shape
214
215    return results, shape