From 360f84f954a2ac2739aabf0168bdefb7b6bedcbf Mon Sep 17 00:00:00 2001
From: Jim Martens <github@2martens.de>
Date: Tue, 2 Jul 2019 11:13:38 +0200
Subject: [PATCH] Added visualise task

Signed-off-by: Jim Martens <github@2martens.de>
---
 setup.py                            |  2 +-
 src/twomartens/masterthesis/cli.py  | 54 +++++++++++++++++++++++++++++
 src/twomartens/masterthesis/main.py | 11 ++++++
 3 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 5ddb0e5..d09101f 100644
--- a/setup.py
+++ b/setup.py
@@ -39,7 +39,7 @@ setup(
     },
     python_requires="~=3.6",
     install_requires=["tensorflow-gpu", "Pillow", "h5py", "numpy", "opencv-python", "scikit-learn", "tqdm",
-                      "beautifulsoup4", "matplotlib", "protobuf", "imutils"],
+                      "beautifulsoup4", "matplotlib", "protobuf", "imutils", "matplotlib"],
     license="Apache License 2.0",
     classifiers=[
         "Operating System :: OS Independent",
diff --git a/src/twomartens/masterthesis/cli.py b/src/twomartens/masterthesis/cli.py
index 191cde8..89e00f7 100644
--- a/src/twomartens/masterthesis/cli.py
+++ b/src/twomartens/masterthesis/cli.py
@@ -346,3 +346,57 @@ def prepare(args: argparse.Namespace) -> None:
         pickle.dump(file_names_instances, file)
     with open(f"{args.ground_truth_path}/instances.bin", "wb") as file:
         pickle.dump(instances, file)
+
+
+def visualise(args: argparse.Namespace) -> None:
+    import pickle
+    
+    from matplotlib import pyplot
+    import numpy as np
+    from PIL import Image
+
+    from twomartens.masterthesis.ssd_keras.eval_utils import coco_utils
+    
+    with open(f"{args.ground_truth_path}/photo_paths.bin", "rb") as file:
+        file_names = pickle.load(file)
+    with open(f"{args.ground_truth_path}/instances.bin", "rb") as file:
+        instances = pickle.load(file)
+
+    output_path = f"{args.output_path}/visualise/{args.trajectory}"
+    annotation_file_train = f"{args.coco_path}/annotations/instances_train2014.json"
+    _, _, cats_to_names, _ = coco_utils.get_coco_category_maps(annotation_file_train)
+
+    colors = pyplot.cm.hsv(np.linspace(0, 1, 81)).tolist()
+    classes = ['background'].extend(cats_to_names)
+    
+    i = 0
+    nr_images = len(file_names[args.trajectory])
+    nr_digits = math.ceil(math.log10(nr_images))
+    for file_name, labels in zip(file_names[args.trajectory], instances[args.trajectory]):
+        if not labels:
+            continue
+        
+        # only loop through selected trajectory
+        with Image.open(file_name) as image:
+            figure = pyplot.figure(figsize=(20, 12))
+            pyplot.imshow(image)
+        
+            current_axis = pyplot.gca()
+        
+            for instance in labels:
+                bbox = instance['bbox']
+                # Transform the predicted bounding boxes for the 300x300 image to the original image dimensions.
+                xmin = bbox[0]
+                ymin = bbox[1]
+                xmax = bbox[2]
+                ymax = bbox[3]
+                color = colors[int(instance['coco_id'])]
+                label = f"{classes[int(instance['coco_id'])]}"
+                current_axis.add_patch(
+                    pyplot.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, color=color, fill=False, linewidth=2))
+                current_axis.text(xmin, ymin, label, size='x-large', color='white',
+                                  bbox={'facecolor': color, 'alpha': 1.0})
+            pyplot.savefig(f"{output_path}/{str(i).zfill(nr_digits)}")
+            figure.clear()
+        
+        i += 1
diff --git a/src/twomartens/masterthesis/main.py b/src/twomartens/masterthesis/main.py
index c2d606d..0804c79 100644
--- a/src/twomartens/masterthesis/main.py
+++ b/src/twomartens/masterthesis/main.py
@@ -43,12 +43,14 @@ def main() -> None:
     train_parser = sub_parsers.add_parser("train", help="Train a network")
     evaluate_parser = sub_parsers.add_parser("evaluate", help="Evaluate a network")
     test_parser = sub_parsers.add_parser("test", help="Test a network")
+    visualise_parser = sub_parsers.add_parser("visualise", help="Visualise the ground truth")
     
     # build sub parsers
     _build_prepare(prepare_parser)
     _build_train(train_parser)
     _build_test(test_parser)
     _build_evaluate(evaluate_parser)
+    _build_visualise(visualise_parser)
     
     args = parser.parse_args()
     
@@ -60,6 +62,8 @@ def main() -> None:
         cli.test(args)
     elif args.action == "prepare":
         cli.prepare(args)
+    elif args.actions == "visualise":
+        cli.visualise(args)
 
 
 def _build_prepare(parser: argparse.ArgumentParser) -> None:
@@ -155,5 +159,12 @@ def _build_ssd_evaluate(parser: argparse.ArgumentParser) -> None:
     parser.add_argument("iteration", type=int, help="the validation iteration to use")
 
 
+def _build_visualise(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--coco_path", type=str, help="the path to the COCO data set")
+    parser.add_argument("--ground_truth_path", type=str, help="path to the prepared ground truth directory")
+    parser.add_argument("--output_path", type=str, help="path to the output directory")
+    parser.add_argument("trajectory", type=int, help="trajectory to visualise")
+
+
 if __name__ == "__main__":
     main()