implement automatic patient bed detection

mu_map/data/remove_bed.py

@@ -13,9 +13,64 @@ import numpy as np
 from mu_map.dataset.util import align_images
 
 
+DEFAULT_BED_TRESHOLD = 0.05
 DEFAULT_BED_CONTOURS_FILENAME = "bed_contours.json"
 
 
+def find_bed_contour(
+    mu_map: np.ndarray, threshold: float = DEFAULT_BED_TRESHOLD
+) -> np.ndarray:
+    """
+    Automatically find the contour of the patient bed in a mu map.
+
+    The algorithm works as follows:
+    1. Select the central slice of the mu map and apply a threshold which includes
+       the patient bed but leaves a gap to the patient body. The central slice is
+       selected because some border slices in our mu maps were broken (see mu_map/data/review_mu_map.py).
+       The patient bed looks the same in all slices.
+    2. Find all regions and select the one containing the lowest pixels as the patient bed.
+    3. Dilate the selected region in order to include border pixels with a lower attenuation
+       coefficient than the threshold.
+
+    Parameters
+    ----------
+    mu_map: np.ndarray
+        the mu map in which the contour of the patient bed is searched
+    threshold: float, optional
+        threshold applied to the mu map which selects the patient bed while
+        separating it from the patient body
+
+    Returns
+    -------
+    np.ndarray
+        the contour of the patient bed
+    """
+    # select the central slice
+    mu_map = mu_map[mu_map.shape[0] // 2]
+    # apply the threshold
+    _, mu_map = cv.threshold(mu_map, threshold, 255, cv.THRESH_BINARY)
+    # convert to uint8 image for further processing with OpenCV
+    mu_map = mu_map.astype(np.uint8)
+    # find contours of all regions
+    cnts, _ = cv.findContours(mu_map, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
+    # compute max y coordinate for all contours
+    max_y_coords = list(map(lambda cnt: cnt[:, 0, 1].max(), cnts))
+    # select contour with max y coordinate as the bed contour
+    bed_cnt = cnts[np.argmax(max_y_coords)]
+
+    # dilate the bed contour to catch border pixels
+    _img = np.zeros(mu_map.shape, np.uint8)
+    _img = cv.drawContours(_img, [bed_cnt], -1, 255, -1)
+    kernel = np.ones((3, 3), np.uint8)
+    _img = cv.dilate(_img, kernel, iterations=1)
+
+    # find the dilated contour again
+    cnts, _ = cv.findContours(_img, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
+    bed_cnt = cnts[0]
+
+    return bed_cnt
+
+
 def load_contours(filename: str, as_ndarry: bool = True) -> Dict[int, np.ndarray]:
     """
     Load contours from a JSON file.