Fix unit tests

tests/test_confidence_generator.py

@@ -20,6 +20,13 @@ def generate_traversability_test_signal(N=500, T=10, events=[3, 8], event_length
 
 
 def test_confidence_generator():
+    from wild_visual_navigation.utils.testing import make_results_folder
+    from wild_visual_navigation.visu import get_img_from_fig
+    from os.path import join
+
+    # Create test directory
+    outpath = make_results_folder("test_confidence_generator")
+
     device = torch.device("cpu")
     N = 1000
     sigma_factor = 0.5
@@ -149,7 +156,13 @@ def test_confidence_generator():
     axs[2].set_ylabel("Confidence")
     axs[2].legend(loc="upper right")
 
-    plt.show()
+    img = get_img_from_fig(fig)
+    img.save(
+        join(
+            outpath,
+            "confidence_generator_test.png",
+        )
+    )
 
 
 if __name__ == "__main__":

tests/test_dino_time.py

@@ -1,16 +1,15 @@
-from wild_visual_navigation import WVN_ROOT_DIR
+# from wild_visual_navigation import WVN_ROOT_DIR
 from pytictac import Timer
 from wild_visual_navigation.feature_extractor import (
     DinoInterface,
-    DinoTrtInterface,
     # TrtModel,
 )
 
 # from collections import namedtuple, OrderedDict
 # from torchvision import transforms as T
-import cv2
-import os
+# import os
 import torch
+from wild_visual_navigation.utils.testing import load_test_image, get_dino_transform
 
 # import tensorrt as trt
 # import numpy as np
@@ -19,23 +18,21 @@
 def test_dino_interfacer():
     device = "cuda" if torch.cuda.is_available() else "cpu"
     di = DinoInterface(device)
+    transform = get_dino_transform()
 
-    np_img = cv2.imread(os.path.join(WVN_ROOT_DIR, "assets/images/forest_clean.png"))
-    img = torch.from_numpy(cv2.cvtColor(np_img, cv2.COLOR_BGR2RGB)).to(device)
-    img = img.permute(2, 0, 1)
-    img = (img.type(torch.float32) / 255)[None]
+    img = load_test_image().to(device)
 
     #####################################################################################
     for i in range(5):
         im = img + torch.rand(img.shape, device=img.device) / 100
-        di.inference(di.transform(im))
+        di.inference(transform(im))
 
     #####################################################################################
     with Timer("BS1 Dino Inference: "):
         for i in range(5):
             im = img + torch.rand(img.shape, device=img.device) / 100
             with Timer("BS1 Dino Single: "):
-                di.inference(di.transform(im))
+                di.inference(transform(im))
 
     #####################################################################################
     # img = img.repeat(4, 1, 1, 1)
@@ -46,14 +43,15 @@ def test_dino_interfacer():
     #             res = di.inference(di.transform(im))
 
     #####################################################################################
-    # Conversion from ONNX model (https://github.com/facebookresearch/dino)
-    exported_trt_file = "dino_exported.trt"
-    exported_trt_path = os.path.join(WVN_ROOT_DIR, "assets/dino", exported_trt_file)
-    di_trt = DinoTrtInterface(exported_trt_path, device)
+    # # Conversion from ONNX model (https://github.com/facebookresearch/dino)
+    # from wild_visual_navigation.feature_extractor import DinoTrtInterface
+    # exported_trt_file = "dino_exported.trt"
+    # exported_trt_path = os.path.join(WVN_ROOT_DIR, "assets/dino", exported_trt_file)
+    # di_trt = DinoTrtInterface(exported_trt_path, device)
 
-    with Timer("TensorRT Inference: "):
-        im = img + torch.rand(img.shape, device=img.device) / 100
-        di_trt.inference(di.transform(im).contiguous())
+    # with Timer("TensorRT Inference: "):
+    #     im = img + torch.rand(img.shape, device=img.device) / 100
+    #     di_trt.inference(di.transform(im).contiguous())
 
     #####################################################################################
     # Conversion using the torch_tensorrt library: https://github.com/pytorch/TensorRT

tests/test_feature_extractor.py

@@ -1,53 +1,51 @@
-from wild_visual_navigation import WVN_ROOT_DIR
 from wild_visual_navigation.feature_extractor import FeatureExtractor
 from wild_visual_navigation.visu import get_img_from_fig
+from wild_visual_navigation.utils.testing import load_test_image, get_dino_transform, make_results_folder
+from os.path import join
+from pytictac import Timer
 import matplotlib.pyplot as plt
 import torch
-from torchvision import transforms as T
-from pathlib import PurePath, Path
-import os
-import cv2
+import itertools
 
 
 def test_feature_extractor():
-    segmentation_type = "none"
-    feature_type = "dino"
-
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    fe = FeatureExtractor(device, segmentation_type=segmentation_type, feature_type=feature_type)
-
-    transform = T.Compose(
-        [
-            T.Resize(448, T.InterpolationMode.NEAREST),
-            T.CenterCrop(448),
-        ]
-    )
-
-    np_img = cv2.imread(os.path.join(WVN_ROOT_DIR, "assets/images/forest_clean.png"))
-    img = torch.from_numpy(cv2.cvtColor(np_img, cv2.COLOR_BGR2RGB)).to(device)
-    img = img.permute(2, 0, 1)
-    img = (img.type(torch.float32) / 255)[None]
-    adj, feat, seg, center = fe.extract(transform(img.clone()))
-
-    p = PurePath(WVN_ROOT_DIR).joinpath(
-        "results",
-        "test_feature_extractor",
-        f"forest_clean_graph_{segmentation_type}.png",
-    )
-    Path(p.parent).mkdir(parents=True, exist_ok=True)
-
-    # Plot result as in colab
-    fig, ax = plt.subplots(1, 2, figsize=(5 * 3, 5))
-
-    ax[0].imshow(transform(img).permute(0, 2, 3, 1)[0].cpu())
-    ax[0].set_title("Image")
-    ax[1].imshow(seg.cpu(), cmap=plt.colormaps.get("inferno"))
-    ax[1].set_title("Segmentation")
-    plt.tight_layout()
-
-    # Store results to test directory
-    img = get_img_from_fig(fig)
-    img.save(str(p))
+    segmentation_types = ["none", "grid", "slic", "random", "stego"]
+    feature_types = ["dino", "dinov2", "stego"]
+    backbone_types = ["vit_small", "vit_base", "vit_small_reg", "vit_base_reg"]
+
+    for seg_type, feat_type, back_type in itertools.product(segmentation_types, feature_types, backbone_types):
+        if seg_type == "stego" and feat_type != "stego":
+            continue
+
+        with Timer(f"Running seg [{seg_type}], feat [{feat_type}], backbone [{back_type}]"):
+            try:
+                fe = FeatureExtractor(
+                    device, segmentation_type=seg_type, feature_type=feat_type, backbone_type=back_type
+                )
+            except Exception:
+                print("Not available")
+                continue
+
+        img = load_test_image().to(device)
+        transform = get_dino_transform()
+        outpath = make_results_folder("test_feature_extractor")
+
+        # Compute
+        edges, feat, seg, center, dense_feat = fe.extract(transform(img.clone()))
+
+        # Plot result as in colab
+        fig, ax = plt.subplots(1, 2, figsize=(5 * 3, 5))
+
+        ax[0].imshow(transform(img).permute(0, 2, 3, 1)[0].cpu())
+        ax[0].set_title("Image")
+        ax[1].imshow(seg.cpu(), cmap=plt.colormaps.get("inferno"))
+        ax[1].set_title("Segmentation")
+        plt.tight_layout()
+
+        # Store results to test directory
+        img = get_img_from_fig(fig)
+        img.save(join(outpath, f"forest_clean_graph_{seg_type}_{feat_type}.png"))
 
 
 if __name__ == "__main__":

tests/test_image_projector.py

@@ -7,24 +7,19 @@ def test_image_projector():
 
 
 def test_supervision_projection():
-    from wild_visual_navigation import WVN_ROOT_DIR
     from wild_visual_navigation.image_projector import ImageProjector
     from wild_visual_navigation.visu import get_img_from_fig
-    from PIL import Image
+    from wild_visual_navigation.utils.testing import load_test_image, make_results_folder
     from liegroups.torch import SE3, SO3
     import matplotlib.pyplot as plt
     import torch
-    import torchvision.transforms as transforms
-    import os
     from os.path import join
     from kornia.utils import tensor_to_image
-    from stego.src import remove_axes
+    from stego.utils import remove_axes
     import random
 
-    to_tensor = transforms.ToTensor()
-
     # Create test directory
-    os.makedirs(join(WVN_ROOT_DIR, "results", "test_image_projector"), exist_ok=True)
+    outpath = make_results_folder("test_image_projector")
 
     # Define number of cameras (batch)
     B = 100
@@ -36,19 +31,17 @@ def test_supervision_projection():
 
     # Extrisics
     pose_camera_in_world = torch.eye(4)[None]
-
     # Image size
-    H = 1080
-    W = 1440
+    H = torch.tensor(1080)
+    W = torch.tensor(1440)
 
     # Create projector
     im = ImageProjector(K, H, W)
 
     # Load image
-    pil_image = Image.open(join(WVN_ROOT_DIR, "assets/images/forest_clean.png"))
+    torch_image = load_test_image()
 
-    # Convert to torch
-    torch_image = to_tensor(pil_image)
+    # Resize
     torch_image = im.resize_image(torch_image)
     mask = (torch_image * 0.0)[None]
 
@@ -68,7 +61,7 @@ def test_supervision_projection():
         delta = SE3(R_WC, rho).as_matrix()[None]  # Pose matrix of camera in world frame
         pose_base_in_world = pose_base_in_world @ delta
         pose_footprint_in_base = torch.eye(4)[None]
-        print(delta, pose_base_in_world)
+        # print(delta, pose_base_in_world)
 
         twist = torch.rand((3,))
         supervision = torch.rand((10,))
@@ -105,22 +98,15 @@ def test_supervision_projection():
     fig, ax = plt.subplots(1, 2, figsize=(2 * 5, 5))
     ax[0].imshow(tensor_to_image(torch_image))
     ax[0].set_title("Image")
-    ax[1].imshow(tensor_to_image(mask))
+    ax[1].imshow(tensor_to_image(mask[0]))
     ax[1].set_title("Labels")
 
     remove_axes(ax)
     plt.tight_layout()
 
     # Store results to test directory
     img = get_img_from_fig(fig)
-    img.save(
-        join(
-            WVN_ROOT_DIR,
-            "results",
-            "test_image_projector",
-            "forest_clean_supervision_projection.png",
-        )
-    )
+    img.save(join(outpath, "forest_clean_supervision_projection.png"))
 
 
 if __name__ == "__main__":

tests/test_kornia.py

@@ -1,6 +1,5 @@
-from wild_visual_navigation import WVN_ROOT_DIR
 from wild_visual_navigation.visu import get_img_from_fig
-import os
+from wild_visual_navigation.utils.testing import make_results_folder
 from os.path import join
 import matplotlib.pyplot as plt
 import torch
@@ -47,8 +46,8 @@ def test_draw_polygon():
     # Draw
     k_out = draw_convex_polygon(img, poly, color)
     # Show
-    os.makedirs(join(WVN_ROOT_DIR, "results", "test_kornia"), exist_ok=True)
+    outpath = make_results_folder("test_kornia")
     fig, ax = plt.subplots(1, 1, figsize=(5, 5))
     ax.imshow(tensor_to_image(k_out))
     out_img = get_img_from_fig(fig)
-    out_img.save(join(WVN_ROOT_DIR, "results", "test_kornia", "polygon_test.png"))
+    out_img.save(join(outpath, "polygon_test.png"))

tests/test_monitoring.py

@@ -1,5 +1,5 @@
 from wild_visual_navigation.utils import SystemLevelGpuMonitor, accumulate_memory
-from pytictac import SystemLevelTimer, accumulate_time
+from pytictac import accumulate_time
 import time
 import torch
 
@@ -15,13 +15,13 @@ def __init__(self):
         @accumulate_time
         def test_memory_then_timing(self, s):
             time.sleep(s / 1000)
-            self.tensors.append(torch.zeros((s, s), device="cuda"))
+            self.tensors.append(torch.zeros((10 * s, 10 * s), device="cuda"))
 
         @accumulate_time
         @accumulate_memory
         def test_timing_then_memory(self, s):
             time.sleep(s / 1000)
-            self.tensors.append(torch.zeros((4, s, s), device="cuda"))
+            self.tensors.append(torch.zeros((4, 10 * s, 10 * s), device="cuda"))
 
     # Create objects
     my_test = MyTest()
@@ -33,13 +33,13 @@ def test_timing_then_memory(self, s):
         store_samples=True,
         skip_n_samples=1,
     )
-    time_monitor = SystemLevelTimer(
-        objects=[my_test],
-        names=["test"],
-    )
+    # time_monitor = ClassContextTimer(
+    #     objects=[my_test],
+    #     names=["test"],
+    # )
 
     # Run loop
-    for n in range(400):
+    for n in range(100):
         print(f"step {n}")
         step = n
         t = n / 10
@@ -48,7 +48,7 @@ def test_timing_then_memory(self, s):
         my_test.test_timing_then_memory(n)
 
     gpu_monitor.store("/tmp")
-    time_monitor.store("/tmp")
+    # time_monitor.store("/tmp")
 
 
 if __name__ == "__main__":

tests/test_optical_flow.py

@@ -1,15 +1,22 @@
-from pytorch_pwc.network import PwcFlowEstimator
-from pytorch_pwc import PWC_ROOT_DIR
-import torch
-from wild_visual_navigation import WVN_ROOT_DIR
-from wild_visual_navigation.visu import LearningVisualizer
-import numpy as np
-import PIL
-import os
-from pytictac import Timer
+import pytest
+import sys
 
-if __name__ == "__main__":
+try:
+    from pytorch_pwc.network import PwcFlowEstimator
+    from pytorch_pwc import PWC_ROOT_DIR
+except ImportError:
+    pass
+
+
+@pytest.mark.skipif("pytorch_pwc" not in sys.modules, reason="requires the pytorch_pwc library")
+def pytorch_pwc_test():
     import os
+    import torch
+    from wild_visual_navigation import WVN_ROOT_DIR
+    from wild_visual_navigation.visu import LearningVisualizer
+    import numpy as np
+    import PIL
+    from pytictac import Timer
 
     tenOne = torch.FloatTensor(
         np.ascontiguousarray(
@@ -39,3 +46,7 @@
     visu = LearningVisualizer(p_visu=os.path.join(WVN_ROOT_DIR, "results/test_visu"), store=True)
     visu.plot_optical_flow(res, tenOne, tenTwo)
     print("done")
+
+
+if __name__ == "__main__":
+    pytorch_pwc_test()