Plotting stuff2 (#13)

* wavemap: Add example pipeline start

* wavemap: Running

* wavemap: save progress

* validation: Fix calibration slightly

* wavemap: Correct pose direction

* validation: visualize PCs

* validation: Remove snr filter

* validation: Handle empty radar data

* validation: Debug

* validation: Work on filtering

* validation: Fix depth dict

* plotting: more stuff

radarmeetsvision/metric_depth_network/dataset/blearndataset.py

@@ -49,12 +49,10 @@ def __init__(self, dataset_dir, mode, size, index_min=0, index_max=-1, depth_pri
 
         self.rgb_dir = self.dir / "rgb"
         self.rgb_template = "{:05d}_rgb.jpg"
-        self.rgb_template_alt = "{:04d}_rgb.jpg"
         self.rgb_mask = r'([0-9]*)_rgb.jpg'
 
         self.depth_dir = self.dir / "depth"
         self.depth_template = "{:05d}_d.npy"
-        self.depth_template_alt = "{:04d}_d.npy"
         self.depth_normalized_template = "{:05d}_dn.npy"
         self.depth_min, self.depth_max, self.depth_range = self.get_depth_range()
 
@@ -114,10 +112,6 @@ def __getitem__(self, item):
         index = int(self.filelist[item])
         img_path = self.rgb_dir / self.rgb_template.format(index)
 
-        # TODO: Fix properly, screwed up the hypersim dataset with :04d
-        if not img_path.is_file():
-            img_path = self.rgb_dir / self.rgb_template_alt.format(index)
-
         image_cv = cv2.imread(str(img_path))
         image = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB) / 255.0
 
@@ -152,15 +146,10 @@ def __len__(self):
     def get_depth(self, index):
         depth = None
         depth_path = self.depth_dir / self.depth_template.format(index)
-        depth_path_alt = self.depth_dir / self.depth_template_alt.format(index)
         depth_normalized_path = self.depth_dir / self.depth_normalized_template.format(index)
         if depth_path.is_file():
             depth = np.load(depth_path)
 
-        # TODO: Fix properly, screwed up the hypersim dataset with :04d
-        elif depth_path_alt.is_file():
-            depth = np.load(depth_path_alt)
-
         elif depth_normalized_path.is_file():
             if self.depth_range is not None and self.depth_min is not None:
                 depth_normalized = np.load(depth_normalized_path)
@@ -240,7 +229,7 @@ def get_depth_prior(self, index, img_copy, depth):
         if self.depth_prior_dir.is_dir():
             depth_prior = np.load(str(self.depth_prior_dir / self.depth_prior_template.format(index))).astype(np.float32)
             # TODO: Better detecting if dataset is normalized or not
-            if not (depth_prior > 0.0).any() and self.depth_min is not None and self.depth_range is not None:
+            if (depth_prior > 0.0).any() and depth_prior.max() <= 1.0 and self.depth_min is not None and self.depth_range is not None:
                 depth_prior_valid_mask = (depth_prior > 0.0) & (depth_prior <= 1.0)
                 depth_prior[depth_prior_valid_mask] *= self.depth_range + self.depth_min
 
@@ -256,29 +245,42 @@ def compute_depth_prior(self, img_copy, depth):
         number_of_points = 50
         radius_prior_pixel = 5
         maximum_number_of_corners = 5
-        height, width = img_copy.shape[0], img_copy.shape[1]
+        height, width = img_copy.shape[:2]
 
         circular_mask = self.create_circular_mask(2 * height, 2 * width, radius=radius_prior_pixel)
         depth_prior = np.zeros((height, width), dtype=np.float32)
         img = cv2.cvtColor(img_copy, cv2.COLOR_BGR2GRAY)
         corners = cv2.goodFeaturesToTrack(img, number_of_points, quality_level, minimum_distance)
 
+        def remove_duplicates(corners, threshold=10.0):
+            """Remove duplicate corners based on a threshold distance."""
+            unique_corners = []
+            for corner in corners:
+                x, y = corner
+                if all(np.linalg.norm(np.array([x, y]) - np.array([ux, uy])) > threshold for ux, uy in unique_corners):
+                    unique_corners.append((x, y))
+            return unique_corners
+
         if corners is not None:
             corners = corners.squeeze()
+            corners = remove_duplicates(corners)
+
+            # Shuffle and limit the number of corners
             random.shuffle(corners)
             num_corners_to_use = random.randint(1, min(len(corners), maximum_number_of_corners))
             good_corners = 0
 
             # Process the selected corners
-            for c in corners:
+            for c in corners[:num_corners_to_use]:
                 if isinstance(c, np.ndarray):
                     x, y = c
-                    translated_mask = circular_mask[int(height - y):int(2 * height - y), int(width - x):int(2 * width - x)]
+                    translated_mask = circular_mask[int(height - y):int(2 * height - y), 
+                                                    int(width - x):int(2 * width - x)]
                     depth_at_prior = np.mean(depth_prior[translated_mask & ~np.isnan(depth_prior)])
                     if depth_at_prior == 0.0:
                         depth_at_corner = np.mean(depth[translated_mask & ~np.isnan(depth)])
 
-                        if depth_at_corner > self.depth_min and depth_at_corner < self.depth_max:
+                        if self.depth_min < depth_at_corner < self.depth_max:
                             depth_prior += depth_at_corner * translated_mask
                             good_corners += 1
 

scripts/evaluation/create_scatter_plot_split.py

@@ -0,0 +1,184 @@
+import argparse
+from pathlib import Path
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.gridspec import GridSpec
+from PIL import Image
+
+# Configure matplotlib to match the example style
+import matplotlib as mpl
+mpl.rcParams['font.size'] = 8
+mpl.rcParams['figure.figsize'] = [9, 2.7]
+mpl.rcParams['lines.linewidth'] = 0.6
+mpl.rcParams['grid.linewidth'] = 0.5
+mpl.rcParams['axes.linewidth'] = 0.6
+mpl.rcParams['axes.labelsize'] = 8
+mpl.rcParams['axes.titlesize'] = 8
+mpl.rcParams['legend.fontsize'] = 8
+mpl.rcParams['xtick.labelsize'] = 8
+mpl.rcParams['ytick.labelsize'] = 8
+mpl.rcParams['text.usetex'] = True
+mpl.rcParams['font.family'] = 'serif'
+
+def unnormalize_depth(depth, depth_min=0.19983673095703125, depth_max=120.49285888671875):
+    """
+    Unnormalizes a depth file given min and max depth.
+    Assumes input is in range [0, 1] for valid (nonzero) values.
+    Zeros remain unchanged as they encode no depth.
+    """
+    valid_mask = depth > 0  # Only unnormalize nonzero values
+    unnormalized = np.zeros_like(depth)  # Start with zeros
+    unnormalized[valid_mask] = depth[valid_mask] * (depth_max - depth_min) + depth_min
+    return unnormalized
+
+
+def load_npy_files(base_dir):
+    # Define directories based on the base directory
+    depth_dir = base_dir / "depth"
+    prediction_rgb_dir = base_dir / "prediction_rgb"
+    prediction_radar_dir = base_dir / "prediction_radar"
+
+    # Check if directories exist
+    if not all(d.is_dir() for d in [depth_dir, prediction_rgb_dir, prediction_radar_dir]):
+        print("One or more directories (depth, prediction_rgb, prediction_radar) are missing.")
+        return {}, {}, {}
+
+    # Gather depth files
+    depth_files = sorted(depth_dir.glob("*_d*.npy"))  # Match both *_d.npy and *_dn.npy
+
+    # Initialize storage for loaded data
+    depth_data = {}
+    rgb_data = {}
+    radar_data = {}
+
+    for depth_file in depth_files:
+        index = depth_file.stem.split("_")[0]
+        rgb_file = prediction_rgb_dir / f"{index}_p.npy"
+        radar_file = prediction_radar_dir / f"{index}_p.npy"
+
+        # Check if corresponding files exist in both prediction directories
+        if rgb_file.exists() and radar_file.exists():
+            try:
+                # Load the depth file and check for unnormalization
+                depth = np.load(depth_file)
+                if depth_file.stem.endswith("_dn"):
+                    depth = unnormalize_depth(depth)
+
+                # Load predictions
+                rgb_data[index] = np.load(rgb_file)
+                radar_data[index] = np.load(radar_file)
+                depth_data[index] = depth
+            except Exception as e:
+                print(f"Error loading files for index {index}: {e}")
+        else:
+            print(f"Missing prediction files for index {index} (RGB or Radar).")
+
+    print(f"Loaded {len(depth_data)} sets of files.")
+    return depth_data, rgb_data, radar_data
+
+
+def calculate_metrics_for_index(depth, rgb_pred, radar_pred):
+    # Mask for valid depth pixels (non-zero depth values)
+    valid_mask = depth > 0
+
+    if np.any(valid_mask):
+        # Calculate average scene depth
+        avg_depth = np.mean(depth[valid_mask])
+
+        # Calculate MARE for RGB predictions
+        rgb_error = np.abs((rgb_pred[valid_mask] - depth[valid_mask]) / depth[valid_mask])
+        mare_rgb = np.mean(rgb_error)
+
+        # Calculate MARE for Radar predictions
+        radar_error = np.abs((radar_pred[valid_mask] - depth[valid_mask]) / depth[valid_mask])
+        mare_radar = np.mean(radar_error)
+
+        return avg_depth, mare_rgb, mare_radar
+    return None, None, None
+
+
+def plot_metrics_step(avg_scene_depths, mare_rgb, mare_radar, output_dir, step):
+    # Create the output directory if it doesn't exist
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Create the scatter plot
+    fig = plt.figure(figsize=(5.4, 2.7))  # Match the style and layout of the example
+    gs = GridSpec(1, 1)  # Single grid for the scatter plot
+
+    ax = fig.add_subplot(gs[0])
+
+    # Scatter plots
+    ax.scatter(avg_scene_depths, mare_rgb, label="RGB-Method-B", color="red", alpha=0.7, s=25, marker='o')
+    ax.scatter(avg_scene_depths, mare_radar, label="Ours-B", color="blue", alpha=0.7, s=25, marker='o')
+
+    # Axis labels, grid, and legend
+    ax.set_xlabel("Average Scene Depth [m]", fontsize=8)
+    ax.set_ylabel("Absolute Relative Error [ ]", fontsize=8)
+    ax.grid(linewidth=0.5)
+    ax.legend(loc="upper right", fontsize=8)
+
+    plt.tight_layout()
+
+    # Save the figure for this step
+    output_file = output_dir / f"{step:03d}.png"
+    plt.savefig(str(output_file), bbox_inches="tight", dpi=400)
+    plt.close()
+
+
+def main():
+    # Set up argument parser
+    parser = argparse.ArgumentParser(
+        description="Load corresponding depth, prediction RGB, and prediction radar .npy files, and create step-by-step scatter plots."
+    )
+    parser.add_argument(
+        "base_dir",
+        type=str,
+        help="Path to the base directory containing 'depth', 'prediction_rgb', and 'prediction_radar' subdirectories.",
+    )
+
+    # Parse arguments
+    args = parser.parse_args()
+
+    # Resolve the base directory and create the output directory
+    base_dir = Path(args.base_dir).resolve()
+    output_dir = base_dir / "mare"  # Create 'mare' directory inside the base directory
+
+    # Validate base directory
+    if not base_dir.is_dir():
+        print(f"Provided base_dir '{base_dir}' is not a valid directory.")
+        return
+
+    # Load files
+    depth_data, rgb_data, radar_data = load_npy_files(base_dir)
+
+    # Initialize storage for metrics
+    avg_scene_depths = []
+    mare_rgb = []
+    mare_radar = []
+
+    # Process files in ascending order of indices
+    sorted_indices = sorted(depth_data.keys(), key=lambda x: int(x))  # Sort indices numerically
+    for step, index in enumerate(sorted_indices):
+        depth = depth_data[index]
+        rgb_pred = rgb_data.get(index)
+        radar_pred = radar_data.get(index)
+
+        if rgb_pred is None or radar_pred is None:
+            continue
+
+        # Calculate metrics for the current index
+        avg_depth, rgb_mare, radar_mare = calculate_metrics_for_index(depth, rgb_pred, radar_pred)
+
+        if avg_depth is not None:
+            avg_scene_depths.append(avg_depth)
+            mare_rgb.append(rgb_mare)
+            mare_radar.append(radar_mare)
+
+            # Plot and save the current step
+            plot_metrics_step(avg_scene_depths, mare_rgb, mare_radar, output_dir, step)
+
+    print(f"Step-by-step plots saved in '{output_dir}'.")
+
+
+if __name__ == "__main__":
+    main()

scripts/plotting/plot_single_dataset.py

@@ -0,0 +1,55 @@
+import argparse
+import matplotlib.pyplot as plt
+import numpy as np
+import cv2
+import radarmeetsvision as rmv
+
+from pathlib import Path
+
+class Plotter:
+    def __init__(self, dataset_dir):
+        dataset_dir = Path(dataset_dir)
+        self.prediction_dir = dataset_dir / args.dir
+        if not self.prediction_dir.is_dir():
+            self.prediction_dir.mkdir(exist_ok=True)
+
+    def callback(self, sample, prediction):
+        index = int(sample['index'])
+        prediction_file = self.prediction_dir / f'{index:05d}_p.jpg'
+        prediction_np_file = self.prediction_dir / f'{index:05d}_p.npy'
+        prediction_np = prediction.cpu().numpy().squeeze()
+        np.save(prediction_np_file, prediction_np)
+        plt.imsave(str(prediction_file), prediction_np, cmap='viridis')
+
+def main(args):
+    rmv.setup_global_logger()
+    interface = rmv.Interface()
+    interface.set_encoder('vitb')
+    depth_min = 0.19983673095703125
+    depth_max = 120.49285888671875
+    interface.set_depth_range((depth_min, depth_max))
+    interface.set_output_channels(args.outputchannels)
+
+    interface.set_size(480, 640)
+    interface.set_batch_size(1)
+    interface.set_criterion()
+    interface.set_use_depth_prior(bool(args.usedepthprior))
+
+    interface.load_model(pretrained_from=args.network)
+    _, loader = interface.get_single_dataset(args.dataset, min_index=args.min, max_index=args.max)
+
+    plotter = Plotter(args.dataset)
+    interface.validate_epoch(0, loader, iteration_callback=plotter.callback)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Purely evalute a network')
+    parser.add_argument('--dataset', type=str, required=True, help='Path to the dataset directory')
+    parser.add_argument('--network', type=str, help='Path to the network file')
+    parser.add_argument('--dir', type=str, default='prediction')
+    parser.add_argument('--min', type=int, default=0)
+    parser.add_argument('--max', type=int, default=-1)
+    parser.add_argument('--outputchannels', type=int, default=2)
+    parser.add_argument('--usedepthprior', type=int, default=1)
+    args = parser.parse_args()
+    main(args)

scripts/validation_dataset/create_validation_dataset.py

@@ -150,9 +150,10 @@ def generate(self):
                         index = intrinsics_dict[timestamp]['index']
 
                         if points_radar_all is None:
-                            points_radar_all = points_inertial[:3, :]
-                            snr_radar_all = snr_radar
-                            outlier_mask_all = outlier_mask
+                            if len(snr_radar):
+                                points_radar_all = points_inertial[:3, :]
+                                snr_radar_all = snr_radar
+                                outlier_mask_all = outlier_mask
 
                         else:
                             points_radar_all = np.hstack((points_radar_all, points_inertial[:3, :]))
@@ -183,7 +184,7 @@ def generate(self):
                                 f.write(' '.join(map(str, pose_flat)))
                                 f.write('\n')
 
-        self.visualize_point_clouds(self.points, points_radar_all, outlier_mask_all)
+        # self.visualize_point_clouds(self.points, points_radar_all, outlier_mask_all)
 
         fig, axs = plt.subplots(3, 2)
         snr_all_np = np.array(self.snr_all)
@@ -309,8 +310,8 @@ def get_depth_prior(self, points_radar_window, snr_radar_window, noise_radar_win
             radar_depth = radar_image_points[2, :].copy()
             radar_image_points /= radar_depth
 
-            point_mask = (radar_depth > 0.0) & (radar_depth <= 100.0)
-            # point_mask = (radar_distances / radar_depth) < 0.05
+            # This is a good compromise to reject noise
+            point_mask = (radar_depth > 0.0) & (radar_depth <= 10.0)
             if point_mask.sum() > 0:
                 u = radar_image_points[0, point_mask]
                 v = radar_image_points[1, point_mask]