add script that allows to remove the bed from attenuation maps

mu_map/data/remove_bed.py

+import json
+from typing import Dict
+
+import numpy as np
+
+
+DEFAULT_BED_CONTOURS_FILENAME = "bed_contours.json"
+
+
+def load_contours(filename: str) -> Dict[int, np.ndarray]:
+    """
+    Load contours from a json file.
+    The structure of the file is a dict where the key is the id of the according
+    image and the value is a numpy array of the contour.
+
+    :param filename: filename of a json file containing contours
+    :return: a dict mapping ids to contours
+    """
+    with open(filename, mode="r") as f:
+        contours = json.load(f)
+
+        for key, contour in contours.items():
+            del contours[key]
+            contours[int(key)] = np.array(contour).astype(int)
+    return contours
+
+
+if __name__ == "__main__":
+    import argparse
+    from enum import Enum
+    import os
+
+    import cv2 as cv
+
+    from mu_map.data.datasets import MuMapDataset
+    from mu_map.util import to_grayscale, COLOR_BLACK, COLOR_WHITE
+
+    parser = argparse.ArgumentParser(
+        description="draw and save contours to exclude the bed from mu maps",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "dataset_dir", type=str, help="the directory containing the dataset"
+    )
+    parser.add_argument(
+        "--output_file",
+        type=str,
+        default=DEFAULT_BED_CONTOURS_FILENAME,
+        help="default file in dataset dir where the drawn contours are stored",
+    )
+    args = parser.parse_args()
+    args.output_file = os.path.join(args.dataset_dir, args.output_file)
+
+    controls = """
+    Controls:
+    Left click to add points to the current contour.
+
+    q: exit
+    d: delete last point
+    n: save contour and go to the next image
+    v: change the visual mode between drawing contours and hiding the are within
+    """
+    print(controls)
+    print()
+
+    # set bed contours file to None so that existing contours are not used
+    dataset = MuMapDataset(args.dataset_dir, bed_contours_file=None)
+
+    # TODO: implement that existing contours are loaded so that labeling can be continued?
+    bed_contours = {}
+
+    class VisualMode(Enum):
+        DRAW_CONTOURS = 1
+        HIDE_BED = 2
+
+    window_name = "Bed Removal"
+    for i, (_, mu_map) in enumerate(dataset):
+        print(f"Image {str(i + 1):>{len(str(len(dataset)))}}/{len(dataset)}", end="\r")
+        # select the center slice for display (the bed location is constant over all slices)
+        mu_map = mu_map[mu_map.shape[0] // 2]
+
+        # save the points of the contour in a list and defined a mouse callback
+        points = []
+
+        def mouse_callback(event, x, y, flags, param):
+            if event == cv.EVENT_LBUTTONUP:
+                points.append((x, y))
+
+        # create a window for display
+        cv.namedWindow(window_name, cv.WINDOW_NORMAL)
+        cv.resizeWindow(window_name, 1024, 1024)
+        cv.setMouseCallback(window_name, mouse_callback)
+
+        # set initial visual mode
+        visual_mode = VisualMode.DRAW_CONTOURS
+        while True:
+            # compute image to display
+            to_show = to_grayscale(mu_map)
+
+            if visual_mode == VisualMode.DRAW_CONTOURS:
+                # draw lines between all points
+                for p1, p2 in zip(points[:-1], points[1:]):
+                    to_show = cv.line(to_show, p1, p2, color=COLOR_WHITE, thickness=1)
+                # close the contour
+                if len(points) > 0:
+                    to_show = cv.line(
+                        to_show, points[0], points[-1], color=COLOR_WHITE, thickness=1
+                    )
+
+                # draw all points as circles
+                for point in points:
+                    to_show = cv.circle(
+                        to_show, point, radius=2, color=COLOR_BLACK, thickness=-1
+                    )
+                    to_show = cv.circle(
+                        to_show, point, radius=2, color=COLOR_WHITE, thickness=1
+                    )
+            else:
+                # eliminate area inside the contour
+                _points = np.array(points).astype(int)
+                to_show = cv.drawContours(
+                    to_show, [_points], -1, COLOR_BLACK, thickness=-1
+                )
+
+            # visualize image and handle inputs
+            cv.imshow(window_name, to_show)
+            key = cv.waitKey(100)
+            if key == ord("q"):
+                exit(0)
+            elif key == ord("d"):
+                points = points[:-1]
+            elif key == ord("n"):
+                break
+            elif key == ord("v"):
+                visual_mode = (
+                    VisualMode.DRAW_CONTOURS
+                    if visual_mode == VisualMode.HIDE_BED
+                    else VisualMode.HIDE_BED
+                )
+
+        # remove current window
+        cv.destroyWindow(window_name)
+
+        # save current contour in dict
+        bed_contours[int(dataset.table.loc[i, "id"])] = points
+
+    # write contours to output file
+    with open(args.output_file, mode="w") as f:
+        f.write(json.dumps(data, indent=2, sort_keys=True))