fix remove bed contours script

mu_map/data/remove_bed.py

@@ -7,18 +7,22 @@ import numpy as np
 DEFAULT_BED_CONTOURS_FILENAME = "bed_contours.json"
 
 
-def load_contours(filename: str) -> Dict[int, np.ndarray]:
+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.
 
     :param filename: filename of a json file containing contours
+    :param as_ndarry: directly parse contours as numpy arrays
     :return: a dict mapping ids to contours
     """
     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)
 
@@ -65,7 +69,14 @@ if __name__ == "__main__":
     dataset = MuMapDataset(args.dataset_dir, bed_contours_file=None)
 
     # TODO: implement that existing contours are loaded so that labeling can be continued?
-    bed_contours = {}
+    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 = {}
+    else:
+        bed_contours = {}
 
     class VisualMode(Enum):
         DRAW_CONTOURS = 1
@@ -73,6 +84,12 @@ if __name__ == "__main__":
 
     window_name = "Bed Removal"
     for i, (_, mu_map) in enumerate(dataset):
+        _id = str(int(dataset.table.loc[i, "id"]))
+
+        if _id in bed_contours:
+            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]
@@ -124,6 +141,9 @@ if __name__ == "__main__":
             cv.imshow(window_name, to_show)
             key = cv.waitKey(100)
             if key == ord("q"):
+                # write current contours to output file
+                with open(args.output_file, mode="w") as f:
+                    f.write(json.dumps(bed_contours, indent=2, sort_keys=True))
                 exit(0)
             elif key == ord("d"):
                 points = points[:-1]
@@ -140,8 +160,8 @@ if __name__ == "__main__":
         cv.destroyWindow(window_name)
 
         # save current contour in dict
-        bed_contours[int(dataset.table.loc[i, "id"])] = points
+        bed_contours[_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))
+        f.write(json.dumps(bed_contours, indent=2, sort_keys=True))