using mse loss and distance transform (#73)

Author: mhubii

README.md

@@ -157,6 +157,7 @@ The camera swarm optimization can serve for finding an initial guess to [Monocul
 
 ```shell
 rr-cam-swarm \
+    --collision-meshes \
     --n-cameras 1000 \
     --min-distance 0.5 \
     --max-distance 3.0 \

roboreg/cli/rr_cam_swarm.py

@@ -1,20 +1,17 @@
 import argparse
 import os
-from typing import Tuple
 
 import cv2
 import numpy as np
-import rich
-import rich.progress
 import torch
 
 from roboreg import differentiable as rrd
-from roboreg.io import find_files, parse_camera_info
+from roboreg.io import find_files, parse_camera_info, parse_mono_data
 from roboreg.losses import soft_dice_loss
 from roboreg.optim import LinearParticleSwarm, ParticleSwarmOptimizer
 from roboreg.util import (
     look_at_from_angle,
-    mask_exponential_distance_transform,
+    mask_exponential_decay,
     overlay_mask,
     random_fov_eye_space_coordinates,
 )
@@ -173,58 +170,6 @@ def args_factory() -> argparse.Namespace:
     return parser.parse_args()
 
 
-def parse_data(
-    path: str,
-    image_pattern: str,
-    mask_pattern: str,
-    joint_states_pattern: str,
-    n_samples: int = 5,
-    device: str = "cuda",
-) -> Tuple[np.ndarray, torch.Tensor, torch.Tensor]:
-    image_files = find_files(path, image_pattern)
-    mask_files = find_files(path, mask_pattern)
-    joint_states_files = find_files(path, joint_states_pattern)
-
-    rich.print("Found the following files:")
-    rich.print(f"Images: {image_files}")
-    rich.print(f"Masks: {mask_files}")
-    rich.print(f"Joint states: {joint_states_files}")
-
-    # randomly sample n_samples
-    if n_samples > len(image_files):
-        n_samples = len(image_files)
-    random_indices = np.random.choice(len(image_files), n_samples, replace=False)
-    image_files = np.array(image_files)[random_indices].tolist()
-    mask_files = np.array(mask_files)[random_indices].tolist()
-    joint_states_files = np.array(joint_states_files)[random_indices].tolist()
-
-    rich.print(f"Randomly sampled the following {n_samples} files:")
-    rich.print(f"Images: {image_files}")
-    rich.print(f"Masks: {mask_files}")
-    rich.print(f"Joint states: {joint_states_files}")
-
-    if len(mask_files) != len(joint_states_files):
-        raise ValueError("Number of masks and joint states do not match.")
-
-    images = [
-        cv2.imread(os.path.join(path, file), cv2.IMREAD_COLOR) for file in image_files
-    ]
-    masks = [
-        mask_exponential_distance_transform(
-            cv2.imread(os.path.join(path, file), cv2.IMREAD_GRAYSCALE)
-        )
-        for file in mask_files
-    ]
-    joint_states = [np.load(os.path.join(path, file)) for file in joint_states_files]
-
-    masks = torch.tensor(np.array(masks), dtype=torch.float32, device=device)
-    joint_states = torch.tensor(
-        np.array(joint_states), dtype=torch.float32, device=device
-    )
-
-    return images, joint_states, masks
-
-
 def instantiate_particles(
     n_particles: int,
     height: int,
@@ -293,15 +238,30 @@ def main() -> None:
     height, width, intrinsics = parse_camera_info(
         camera_info_file=args.camera_info_file
     )
-    images, joint_states, masks = parse_data(
+    image_files = find_files(args.path, args.image_pattern)
+    mask_files = find_files(args.path, args.mask_pattern)
+    joint_states_files = find_files(args.path, args.joint_states_pattern)
+    n_samples = args.n_samples
+    if n_samples > len(image_files):  # randomly sample n_samples
+        n_samples = len(image_files)
+    random_indices = np.random.choice(len(image_files), n_samples, replace=False)
+    image_files = np.array(image_files)[random_indices].tolist()
+    mask_files = np.array(mask_files)[random_indices].tolist()
+    joint_states_files = np.array(joint_states_files)[random_indices].tolist()
+    images, joint_states, masks = parse_mono_data(
         path=args.path,
-        image_pattern=args.image_pattern,
-        mask_pattern=args.mask_pattern,
-        joint_states_pattern=args.joint_states_pattern,
-        n_samples=args.n_samples,
-        device=device,
+        image_files=image_files,
+        mask_files=mask_files,
+        joint_states_files=joint_states_files,
+    )
+
+    # pre-process data
+    joint_states = torch.tensor(
+        np.array(joint_states), dtype=torch.float32, device=device
     )
     n_joint_states = joint_states.shape[0]
+    masks = [mask_exponential_decay(mask) for mask in masks]
+    masks = torch.tensor(np.array(masks), dtype=torch.float32, device=device)
 
     # scale image data (memory reduction)
     height = int(height * args.scale)

roboreg/cli/rr_mono_dr.py

@@ -9,9 +9,8 @@
 import rich.progress
 import torch
 
-from roboreg.io import find_files, parse_mono_dr_data
-from roboreg.losses import soft_dice_loss
-from roboreg.util import mask_exponential_distance_transform, overlay_mask
+from roboreg.io import find_files, parse_mono_data
+from roboreg.util import mask_distance_transform, overlay_mask
 from roboreg.util.factories import create_robot_scene, create_virtual_camera
 
 
@@ -49,12 +48,6 @@ def args_factory() -> argparse.Namespace:
         default=1.0,
         help="Gamma for the learning rate scheduler.",
     )
-    parser.add_argument(
-        "--sigma",
-        type=float,
-        default=2.0,
-        help="Sigma for the exponential distance transform on target masks.",
-    )
     parser.add_argument(
         "--display-progress",
         action="store_true",
@@ -144,7 +137,7 @@ def main() -> None:
     image_files = find_files(args.path, args.image_pattern)
     joint_states_files = find_files(args.path, args.joint_states_pattern)
     mask_files = find_files(args.path, args.mask_pattern)
-    images, joint_states, masks = parse_mono_dr_data(
+    images, joint_states, masks = parse_mono_data(
         path=args.path,
         image_files=image_files,
         joint_states_files=joint_states_files,
@@ -155,12 +148,10 @@ def main() -> None:
     joint_states = torch.tensor(
         np.array(joint_states), dtype=torch.float32, device=device
     )
-    masks = [
-        mask_exponential_distance_transform(mask, sigma=args.sigma) for mask in masks
-    ]
-    masks = torch.tensor(np.array(masks), dtype=torch.float32, device=device).unsqueeze(
-        -1
-    )
+    distance_maps = [mask_distance_transform(mask) for mask in masks]
+    distance_maps = torch.tensor(
+        np.array(distance_maps), dtype=torch.float32, device=device
+    ).unsqueeze(-1)
 
     # instantiate camera with default identity extrinsics because we optimize for robot pose instead
     camera = {
@@ -213,7 +204,7 @@ def main() -> None:
         renders = {
             "camera": scene.observe_from("camera"),
         }
-        loss = soft_dice_loss(renders["camera"], masks).mean()
+        loss = torch.nn.functional.mse_loss(distance_maps, renders["camera"])
         optimizer.zero_grad()
         loss.backward()
         optimizer.step()
@@ -240,15 +231,15 @@ def main() -> None:
             # difference left / right render / mask
             difference = (
                 cv2.cvtColor(
-                    np.abs(render - masks[0].squeeze().cpu().numpy()),
+                    np.abs(render - masks[0].astype(np.float32) / 255.0),
                     cv2.COLOR_GRAY2BGR,
                 )
                 * 255.0
             ).astype(np.uint8)
             # overlay segmentation mask
             segmentation_overlay = overlay_mask(
                 image,
-                (masks[0].squeeze().cpu().numpy() * 255.0).astype(np.uint8),
+                masks[0],
                 mode="b",
                 scale=1.0,
             )
@@ -277,7 +268,7 @@ def main() -> None:
     for i, render in enumerate(renders):
         render = render.squeeze().cpu().numpy()
         overlay = overlay_mask(images[i], (render * 255.0).astype(np.uint8), scale=1.0)
-        difference = np.abs(render - masks[i].squeeze().cpu().numpy())
+        difference = np.abs(render - masks[i].astype(np.float32) / 255.0)
 
         cv2.imwrite(os.path.join(args.path, f"dr_overlay_{i}.png"), overlay)
         cv2.imwrite(

roboreg/cli/rr_stereo_dr.py

@@ -9,9 +9,8 @@
 import rich.progress
 import torch
 
-from roboreg.io import find_files, parse_stereo_dr_data
-from roboreg.losses import soft_dice_loss
-from roboreg.util import mask_exponential_distance_transform, overlay_mask
+from roboreg.io import find_files, parse_stereo_data
+from roboreg.util import mask_distance_transform, overlay_mask
 from roboreg.util.factories import create_robot_scene, create_virtual_camera
 
 
@@ -49,12 +48,6 @@ def args_factory() -> argparse.Namespace:
         default=1.0,
         help="Gamma for the learning rate scheduler.",
     )
-    parser.add_argument(
-        "--sigma",
-        type=float,
-        default=2.0,
-        help="Sigma for the exponential distance transform on target masks.",
-    )
     parser.add_argument(
         "--display-progress",
         action="store_true",
@@ -177,7 +170,7 @@ def main() -> None:
     left_mask_files = find_files(args.path, args.left_mask_pattern)
     right_mask_files = find_files(args.path, args.right_mask_pattern)
     left_images, right_images, joint_states, left_masks, right_masks = (
-        parse_stereo_dr_data(
+        parse_stereo_data(
             path=args.path,
             left_image_files=left_image_files,
             right_image_files=right_image_files,
@@ -191,19 +184,13 @@ def main() -> None:
     joint_states = torch.tensor(
         np.array(joint_states), dtype=torch.float32, device=device
     )
-    left_masks = [
-        mask_exponential_distance_transform(mask, sigma=args.sigma)
-        for mask in left_masks
-    ]
-    right_masks = [
-        mask_exponential_distance_transform(mask, sigma=args.sigma)
-        for mask in right_masks
-    ]
-    left_masks = torch.tensor(
-        np.array(left_masks), dtype=torch.float32, device=device
+    left_distance_maps = [mask_distance_transform(mask) for mask in left_masks]
+    right_distance_maps = [mask_distance_transform(mask) for mask in right_masks]
+    left_distance_maps = torch.tensor(
+        np.array(left_distance_maps), dtype=torch.float32, device=device
     ).unsqueeze(-1)
-    right_masks = torch.tensor(
-        np.array(right_masks), dtype=torch.float32, device=device
+    right_distance_maps = torch.tensor(
+        np.array(right_distance_maps), dtype=torch.float32, device=device
     ).unsqueeze(-1)
 
     # instantiate:
@@ -265,10 +252,9 @@ def main() -> None:
             "left": scene.observe_from("left"),
             "right": scene.observe_from("right"),
         }
-        loss = (
-            soft_dice_loss(renders["left"], left_masks).mean()
-            + soft_dice_loss(renders["right"], right_masks).mean()
-        )
+        loss = torch.nn.functional.mse_loss(
+            left_distance_maps, renders["left"]
+        ) + torch.nn.functional.mse_loss(right_distance_maps, renders["right"])
         optimizer.zero_grad()
         loss.backward()
         optimizer.step()
@@ -309,7 +295,7 @@ def main() -> None:
             differences.append(
                 (
                     cv2.cvtColor(
-                        np.abs(left_render - left_masks[0].squeeze().cpu().numpy()),
+                        np.abs(left_render - left_masks[0].astype(np.float32) / 255.0),
                         cv2.COLOR_GRAY2BGR,
                     )
                     * 255.0
@@ -318,7 +304,9 @@ def main() -> None:
             differences.append(
                 (
                     cv2.cvtColor(
-                        np.abs(right_render - right_masks[0].squeeze().cpu().numpy()),
+                        np.abs(
+                            right_render - right_masks[0].astype(np.float32) / 255.0
+                        ),
                         cv2.COLOR_GRAY2BGR,
                     )
                     * 255.0
@@ -329,15 +317,15 @@ def main() -> None:
             segmentation_overlays.append(
                 overlay_mask(
                     left_image,
-                    (left_masks[0].squeeze().cpu().numpy() * 255.0).astype(np.uint8),
+                    left_masks[0],
                     mode="b",
                     scale=1.0,
                 )
             )
             segmentation_overlays.append(
                 overlay_mask(
                     right_image,
-                    (right_masks[0].squeeze().cpu().numpy() * 255.0).astype(np.uint8),
+                    right_masks[0],
                     mode="b",
                     scale=1.0,
                 )
@@ -378,8 +366,10 @@ def main() -> None:
         right_overlay = overlay_mask(
             right_images[i], (right_render * 255.0).astype(np.uint8), scale=1.0
         )
-        left_difference = np.abs(left_render - left_masks[i].squeeze().cpu().numpy())
-        right_difference = np.abs(right_render - right_masks[i].squeeze().cpu().numpy())
+        left_difference = np.abs(left_render - left_masks[i].astype(np.float32) / 255.0)
+        right_difference = np.abs(
+            right_render - right_masks[i].astype(np.float32) / 255.0
+        )
 
         cv2.imwrite(os.path.join(args.path, f"left_dr_overlay_{i}.png"), left_overlay)
         cv2.imwrite(os.path.join(args.path, f"right_dr_overlay_{i}.png"), right_overlay)

roboreg/io.py

@@ -371,13 +371,13 @@ def parse_hydra_data(
     return joint_states, masks, depths
 
 
-def parse_mono_dr_data(
+def parse_mono_data(
     path: str,
     image_files: List[str],
     joint_states_files: List[str],
     mask_files: List[str],
 ) -> Tuple[List[np.ndarray], List[np.ndarray], List[np.ndarray]]:
-    r"""Parse data for monocular differentiable rendering.
+    r"""Parse monocular data.
 
     Args:
         path (str): Path to the data.
@@ -424,7 +424,7 @@ def parse_mono_dr_data(
     return images, joint_states, masks
 
 
-def parse_stereo_dr_data(
+def parse_stereo_data(
     path: str,
     left_image_files: List[str],
     right_image_files: List[str],
@@ -438,7 +438,7 @@ def parse_stereo_dr_data(
     List[np.ndarray],
     List[np.ndarray],
 ]:
-    r"""Parse data for stereo differentiable rendering.
+    r"""Parse stereo data.
 
     Args:
         path (str): Path to the data.