enhancements to bed removal script

mu_map/data/remove_bed.py

 import json
-from typing import Dict
+from typing import Dict, List
 
 import numpy as np
 
@@ -23,10 +23,18 @@ def load_contours(filename: str, as_ndarry: bool = True) -> Dict[int, np.ndarray
     if not as_ndarry:
         return contours
 
-    _map = map(lambda item: (int(item[0]), np.array(item[1]).astype(int)), contours.items())
+    _map = map(
+        lambda item: (int(item[0]), np.array(item[1]).astype(int)), contours.items()
+    )
     return dict(_map)
 
 
+def scale_points(points: List[List[int]], scale: float):
+    for i in range(len(points)):
+        for j in range(len(points[i])):
+            points[i][j] = round(points[i][j] * scale)
+
+
 if __name__ == "__main__":
     import argparse
     from enum import Enum
@@ -34,7 +42,8 @@ if __name__ == "__main__":
 
     import cv2 as cv
 
-    from mu_map.data.datasets import MuMapDataset
+    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(
@@ -44,6 +53,12 @@ if __name__ == "__main__":
     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,
@@ -59,6 +74,9 @@ if __name__ == "__main__":
 
     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
     """
@@ -68,13 +86,14 @@ if __name__ == "__main__":
     # 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?
     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 = {}
 
@@ -82,53 +101,70 @@ if __name__ == "__main__":
         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"
-    for i, (_, mu_map) in enumerate(dataset):
-        _id = str(int(dataset.table.loc[i, "id"]))
+    window_size = 1024
+    cv.namedWindow(window_name, cv.WINDOW_NORMAL)
+    cv.resizeWindow(window_name, window_size, window_size)
+    cv.setMouseCallback(window_name, mouse_callback)
 
-        if _id in bed_contours:
+    ids = list(dataset.table[headers.id])
+    if args.revise_ids:
+        ids = args.revise_ids
+
+    for _i, _id in enumerate(ids):
+        _, mu_map = dataset.getitem_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
 
-        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 = []
+        if args.revise_ids and str(_id) in bed_contours:
+            points.extend(bed_contours[str(_id)].copy())
 
-        def mouse_callback(event, x, y, flags, param):
-            if event == cv.EVENT_LBUTTONUP:
-                points.append((x, y))
+        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
 
-        # create a window for display
-        cv.namedWindow(window_name, cv.WINDOW_NORMAL)
-        cv.resizeWindow(window_name, 1024, 1024)
-        cv.setMouseCallback(window_name, mouse_callback)
+        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 = to_grayscale(mu_map)
+            to_show = mu_map[_slice]
+            to_show = to_grayscale(to_show)
+            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=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=1
+                        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=2, color=COLOR_BLACK, thickness=-1
+                        to_show, point, radius=4, color=COLOR_BLACK, thickness=-1
                     )
                     to_show = cv.circle(
-                        to_show, point, radius=2, color=COLOR_WHITE, thickness=1
+                        to_show, point, radius=4, color=COLOR_WHITE, thickness=1
                     )
             else:
                 # eliminate area inside the contour
@@ -142,13 +178,21 @@ if __name__ == "__main__":
             key = cv.waitKey(100)
             if key == ord("q"):
                 # write current contours to output file
+                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, indent=2, sort_keys=True))
+                    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
@@ -156,12 +200,11 @@ if __name__ == "__main__":
                     else VisualMode.HIDE_BED
                 )
 
-        # remove current window
-        cv.destroyWindow(window_name)
-
         # save current contour in dict
-        bed_contours[_id] = points
+        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, indent=2, sort_keys=True))
+        # write contours to output file
+        with open(args.output_file, mode="w") as f:
+            f.write(json.dumps(bed_contours, sort_keys=True))