diff --git a/mu_map/data/remove_bed.py b/mu_map/data/remove_bed.py index d002f3f49ce48559e59f89af904b8ba51c4428bf..a71f8c5ebd3b7d9d9599e93130c693d2b281785e 100644 --- a/mu_map/data/remove_bed.py +++ b/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))