retuve.defaults.manual_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"""
 16The Default Segmentaition Method for Hip Ultrasound and Xray.
 17
 18Currently, all the defaults are manual methods.
 19
 20AI methods will be added in the future.
 21"""
 22
 23import time
 24from typing import Dict, List, Tuple, Union
 25
 26import cv2
 27import numpy as np
 28import pydicom
 29from nibabel.nifti1 import Nifti1Image
 30from PIL import Image, ImageOps
 31from radstract.data.colors import LabelColours, get_unique_colours
 32from radstract.data.dicom import convert_dicom_to_images
 33from radstract.data.nifti import convert_nifti_to_image_labels
 34
 35from retuve.classes.seg import SegFrameObjects, SegObject
 36from retuve.hip_us.classes.enums import HipLabelsUS
 37from retuve.hip_xray.classes import HipLabelsXray, LandmarksXRay
 38from retuve.keyphrases.config import Config
 39from retuve.logs import log_timings
 40
 41
 42def manual_predict_us_dcm(
 43    dcm: pydicom.FileDataset,
 44    keyphrase: Union[str, Config],
 45    config: Config = None,
 46    seg: Union[str, Nifti1Image] = None,
 47) -> List[SegFrameObjects]:
 48    """
 49    Manual Segmentation for Hip Ultrasound.
 50
 51    :param dcm: The DICOM File
 52    :param keyphrase: The Keyphrase or Config
 53    :param config: The Config
 54    :param seg: The Segmentation File
 55
 56    :return: The Segmentation Results
 57    """
 58    if not config:
 59        config = Config.get_config(keyphrase)
 60
 61    dicom_images = convert_dicom_to_images(
 62        dcm,
 63        crop_coordinates=config.crop_coordinates,
 64        dicom_type=config.dicom_type,
 65    )
 66
 67    return manual_predict_us(dicom_images, keyphrase, config=config, seg=seg)
 68
 69
 70def manual_predict_us(
 71    images: List[Image.Image],
 72    keyphrase: Union[str, Config],
 73    config: Config = None,
 74    seg: Union[str, Nifti1Image] = None,
 75    seg_idx: int = 0,
 76) -> List[SegFrameObjects]:
 77    """
 78    Manual Segmentation for Hip Ultrasound.
 79
 80    :param images: The Images
 81    :param keyphrase: The Keyphrase or Config
 82    :param config: The Config
 83    :param seg: The Segmentation File
 84    :param seg_idx: The Segmentation Index
 85
 86    :return: The Segmentation Results
 87    """
 88
 89    if not config:
 90        config = Config.get_config(keyphrase)
 91
 92    # seg is a nifti file, open it
 93    if not seg:
 94        raise ValueError("seg file is required")
 95
 96    results, _ = convert_nifti_to_image_labels(
 97        seg, crop_coordinates=config.crop_coordinates
 98    )
 99
100    if seg_idx:
101        results = results[seg_idx:]
102
103    classes = {
104        LabelColours.LABEL1: HipLabelsUS.IlliumAndAcetabulum,
105        LabelColours.LABEL2: HipLabelsUS.FemoralHead,
106        LabelColours.LABEL3: HipLabelsUS.OsIchium,
107    }
108
109    timings = []
110    seg_results = []
111
112    for img, result in zip(images, results):
113        start = time.time()
114        img = np.array(img)
115
116        seg_frame_objects = SegFrameObjects(img)
117
118        # get each colour from the seg
119        unique_colours = get_unique_colours(result)
120
121        # convert image to np array
122        for colour in unique_colours:
123            if colour not in classes:
124                continue
125
126            # convert image to np array
127            result = np.array(result)
128
129            mask = np.all(result == colour, axis=-1)
130            mask = mask.astype(np.uint8) * 255
131
132            # get the bounding box
133            contours, _ = cv2.findContours(
134                mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
135            )
136
137            if not contours:
138                continue
139
140            box = cv2.boundingRect(mask)
141            # just get the bottom left and top right coordinates
142            box = (box[0], box[1], box[0] + box[2], box[1] + box[3])
143
144            # get the points on the outside of the mask
145            points = np.concatenate(contours[0], axis=0)
146
147            # drop 80% of the points
148            points = [
149                (int(x), int(y))
150                for i, (x, y) in enumerate(points)
151                if i % 10 == 0
152            ]
153
154            # convert mask to rgb
155            mask = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGB)
156
157            seg_obj = SegObject(
158                points, classes[colour], mask, box=box, conf=1.0
159            )
160
161            seg_frame_objects.append(seg_obj)
162
163        timings.append(time.time() - start)
164
165        if len(seg_frame_objects) == 0:
166            seg_frame_objects = SegFrameObjects.empty(img=img)
167
168        seg_results.append(seg_frame_objects)
169
170    log_timings(timings, title="Manual Segmentation")
171
172    return seg_results
173
174
175def manual_predict_xray(
176    images: List[Image.Image],
177    keyphrase: Union[str, Nifti1Image],
178    config: Config = None,
179    landmark_list: List[Dict[str, List[int]]] = None,
180    scale_landmarks: bool = False,
181) -> Tuple[List[LandmarksXRay], List[SegFrameObjects]]:
182    """
183    Manual Segmentation for Hip Xray.
184
185    :param images: The Images
186    :param keyphrase: The Keyphrase or Config
187    :param config: The Config
188    :param landmark_list: The Landmark List
189    :param scale_landmarks: Whether to scale the landmarks.
190                            This shoudl be set False if your coordinates
191                            are calculated to the 1024x1024 image resize.
192
193    :return: The Landmarks and Segmentation Results
194    """
195    if not config:
196        config = Config.get_config(keyphrase)
197
198    original_wh = [images[0].width, images[0].height]
199
200    if not landmark_list:
201        raise ValueError("landmark_list is required")
202
203    landmark_results = []
204    seg_results = []
205
206    for image, landmarks in zip(images, landmark_list):
207        A, B = image.size
208
209        if scale_landmarks:
210            for key, (x, y) in landmarks.items():
211                landmarks[key] = (
212                    int(x * A / original_wh[0]),
213                    int(y * B / original_wh[1]),
214                )
215
216        landmarks = LandmarksXRay(**landmarks)
217        landmarks_l = [landmarks.pel_l_o, landmarks.pel_l_i, landmarks.fem_l]
218        landmarks_r = [landmarks.pel_r_o, landmarks.pel_r_i, landmarks.fem_r]
219
220        # construct masks for the triangles
221        mask_l = np.zeros((A, B, 3), dtype=np.uint8)
222        mask_r = np.zeros((A, B, 3), dtype=np.uint8)
223
224        cv2.fillPoly(mask_l, [np.array(landmarks_l)], (255, 255, 255))
225        cv2.fillPoly(mask_r, [np.array(landmarks_r)], (255, 255, 255))
226
227        # and bounding boxes
228        box_l = cv2.boundingRect(np.array(landmarks_l))
229        box_r = cv2.boundingRect(np.array(landmarks_r))
230
231        # get the bottom left and top right coordinates
232        box_l = (box_l[0], box_l[1], box_l[0] + box_l[2], box_l[1] + box_l[3])
233        box_r = (box_r[0], box_r[1], box_r[0] + box_r[2], box_r[1] + box_r[3])
234
235        image = np.array(image)
236
237        seg_objects = [
238            SegObject(
239                landmarks_l,
240                HipLabelsXray.AceTriangle,
241                mask_l,
242                box=box_l,
243                conf=1.0,
244            ),
245            SegObject(
246                landmarks_r,
247                HipLabelsXray.AceTriangle,
248                mask_r,
249                box=box_r,
250                conf=1.0,
251            ),
252        ]
253        seg_result = SegFrameObjects(img=image, seg_objects=seg_objects)
254
255        seg_results.append(seg_result)
256        landmark_results.append(landmarks)
257
258    return landmark_results, seg_results