remove_bed.py

"""
Module containing a GUI to label contours of a bed to be extracted from
an attenuation maps.
Additionally, it contains utility functions to load bed contours, remove or
add the bad.
"""
import json
from typing import Dict, List

import cv2 as cv
import numpy as np

from mu_map.dataset.util import align_images


DEFAULT_BED_CONTOURS_FILENAME = "bed_contours.json"


def load_contours(filename: str, as_ndarry: bool = True) -> 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.

    Parameters
    ----------
    filename: str
        filename of a JSON file containing contours
    as_ndarry: bool
        directly parse contours as numpy arrays

    Parameters
    ----------
    Dict
        a dict mapping ids to contours either as lists of int or np.arrays
    """
    with open(filename, mode="r") as f:
        contours = json.load(f)

    if not as_ndarry:
        return contours

    _map = map(
        lambda item: (int(item[0]), np.array(item[1]).astype(int)), contours.items()
    )
    return dict(_map)


def remove_bed(mu_map: np.ndarray, bed_contour: np.ndarray) -> np.ndarray:
    """
    Remove the bed defined by a contour from all slices.

    Parameters
    ----------
    mu_map: np.ndarray
        the mu_map from which the bed is removed
    bed_contour: np.ndarray
        the contour describing where the bed is found

    Returns
    -------
    np.ndarray
        the mu_map with the bed removed
    """
    _mu_map = mu_map.copy()
    for i in range(_mu_map.shape[0]):
        mu_map[i] = cv.drawContours(_mu_map[i], [bed_contour], -1, 0.0, -1)
    return _mu_map


def add_bed(without_bed: np.ndarray, with_bed: np.ndarray, bed_contour: np.ndarray):
    """
    Add the bed to every slice of a mu_map.

    Parameters
    ----------
    without_bed: np.ndarray
        the mu_map without the bed
    with_bed: np.ndarray
        the mu_map with the bed
    bed_contour: np.ndarray
        the contour defining the location of the bed

    Returns
    -------
    np.ndarray
        the mu_map with the bed added
    """
    with_bed, without_bed = align_images(with_bed, without_bed)

    for _slice in range(with_bed.shape[0]):
        with_bed_i = with_bed[_slice]
        without_bed_i = without_bed[_slice]

        cnt_img = np.zeros(without_bed_i.shape, dtype=np.uint8)
        cnt_img = cv.drawContours(cnt_img, [bed_contour], -1, 255, -1)

        without_bed[_slice] = np.where(cnt_img > 0, with_bed_i, without_bed_i)

    return without_bed


if __name__ == "__main__":

    def scale_points(points: List[List[int]], scale: float):
        """
        Utility function to scale all points in a list of points.
        """
        for i in range(len(points)):
            for j in range(len(points[i])):
                points[i][j] = round(points[i][j] * scale)

    import argparse
    from enum import Enum
    import os

    from mu_map.data.prepare import headers
    from mu_map.dataset.default 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(
        "--revise_ids",
        type=int,
        nargs="*",
        help="only revise the contour of an image withe a specified index",
    )
    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
    c: delete all points
    f: forward (next slice)
    b: backward (previous slice)
    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)

    if os.path.isfile(args.output_file):
        try:
            bed_contours = load_contours(args.output_file, as_ndarry=False)
        except:
            print(f"JSON file {args.output_file} is corrupted! Create a new one.")
            bed_contours = {}
            with open(args.output_file, mode="w") as f:
                f.write(json.dumps(bed_contours, sort_keys=True))
    else:
        bed_contours = {}

    class VisualMode(Enum):
        DRAW_CONTOURS = 1
        HIDE_BED = 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
    window_name = "Bed Removal"
    window_size = 1024
    cv.namedWindow(window_name, cv.WINDOW_NORMAL)
    cv.resizeWindow(window_name, window_size, window_size)
    cv.setMouseCallback(window_name, mouse_callback)

    ids = list(dataset.table[headers.id])
    if args.revise_ids:
        ids = args.revise_ids

    for _i, _id in enumerate(ids):
        _, mu_map = dataset.get_item_by_id(_id)

        if str(_id) in bed_contours and not args.revise_ids:
            print(f"Skip {_id} because file already contains these contours")
            continue

        if args.revise_ids and str(_id) in bed_contours:
            points.extend(bed_contours[str(_id)].copy())

        print(
            f"Image {str(_i + 1):>{len(str(len(ids)))}}/{len(ids)}, ID: {_id:>{len(str(max(ids)))}}",
            end="\r",
        )
        # select the center slice for display (the bed location is constant over all slices)
        mu_map = mu_map.squeeze().numpy()
        _slice = 0

        scale = window_size / mu_map.shape[1]
        scale_points(points, scale)

        # set initial visual mode
        visual_mode = VisualMode.DRAW_CONTOURS
        while True:
            # compute image to display
            to_show = mu_map[_slice]
            to_show = to_grayscale(to_show, min_val=mu_map.min(), max_val=mu_map.max())
            to_show = cv.resize(to_show, (window_size, window_size))

            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=2)
                # close the contour
                if len(points) > 0:
                    to_show = cv.line(
                        to_show, points[0], points[-1], color=COLOR_WHITE, thickness=2
                    )

                # draw all points as circles
                for point in points:
                    to_show = cv.circle(
                        to_show, point, radius=4, color=COLOR_BLACK, thickness=-1
                    )
                    to_show = cv.circle(
                        to_show, point, radius=4, 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"):
                # write current contours to output file
                if len(points) > 0:
                    scale_points(points, 1.0 / scale)
                    bed_contours[str(_id)] = points.copy()
                    with open(args.output_file, mode="w") as f:
                        f.write(json.dumps(bed_contours, sort_keys=True))
                exit(0)
            elif key == ord("d"):
                points = points[:-1]
            elif key == ord("c"):
                points.clear()
            elif key == ord("n"):
                break
            elif key == ord("f"):
                _slice = (_slice + 1) % mu_map.shape[0]
            elif key == ord("b"):
                _slice = (_slice - 1) % mu_map.shape[0]
            elif key == ord("v"):
                visual_mode = (
                    VisualMode.DRAW_CONTOURS
                    if visual_mode == VisualMode.HIDE_BED
                    else VisualMode.HIDE_BED
                )

        # save current contour in dict
        scale_points(points, 1.0 / scale)
        bed_contours[str(_id)] = points.copy()
        points.clear()

        # write contours to output file
        with open(args.output_file, mode="w") as f:
            f.write(json.dumps(bed_contours, sort_keys=True))