import torch import torchvision.transforms.functional as TF from torchvision.transforms import Resize, InterpolationMode from typing import Union import numpy as np AZIMUTH_RESOLUTION = 128 ELEVATION_RESOLUTION = 64 # Depth output resolution: height=64, width=128 DEPTH_TARGET_HEIGHT = 64 DEPTH_TARGET_WIDTH = 128 resize_transform = Resize( size=[ELEVATION_RESOLUTION, AZIMUTH_RESOLUTION], interpolation=InterpolationMode.BILINEAR, antialias=True, ) depth_resize_transform = Resize( size=(DEPTH_TARGET_HEIGHT, DEPTH_TARGET_WIDTH), interpolation=InterpolationMode.BILINEAR, antialias=True, ) def translate_radar(radar_data): """ Applies normalization to radar data after batching from dataloader. Called before passing data into the model. Args: radar_data: Batched radar tensor from dataloader Shape: [B, 64, 8, 2, 256, 2] (batch, doppler, azimuth, elevation, range, channels) - Channel 0: raw amplitude values - Channel 1: phase normalized to [-1, 1] (divided by π) Returns: Processed radar tensor with same shape [B, 64, 8, 2, 256, 2] - Channel 0: sqrt(amplitude * 1e-3) for magnitude normalization - Channel 1: phase * π (converted back to radians [-π, π]) """ radar_mag = radar_data[..., 0] # [B, 64, 8, 2, 256] - Extract raw amplitude radar_phase = radar_data[..., 1] # [B, 64, 8, 2, 256] - Extract normalized phase # Normalize amplitude: scale then sqrt radar_mag_processed = torch.sqrt(radar_mag * 1e-6) # Convert phase back to radians: [-1, 1] -> [-π, π] radar_phase_processed = radar_phase * torch.pi # Stack channels back together: [B, 64, 8, 2, 256, 2] radar_data_translated = torch.stack( [radar_mag_processed, radar_phase_processed], dim=-1 ) return radar_data_translated def resize_depth( depth_map: Union[torch.Tensor, np.ndarray], ) -> Union[torch.Tensor, np.ndarray]: """ Process depth map from dataloader (same pipeline as denoiser/control crop_depth): mm -> meters, clamp [0, 11.2] m, normalize to [0, 1], resize to (64, 128) (h, w). Args: depth_map: Depth in millimeters. Torch or numpy. Shapes: (H, W), (B, H, W), or (B, 1, H, W). Returns: Depth in [0, 1], spatial size (64, 128). Shape [B, 64, 128] for batched input. """ is_numpy = isinstance(depth_map, np.ndarray) if is_numpy: depth_map = torch.from_numpy(depth_map) depth_map = depth_map.float() original_shape = depth_map.shape if depth_map.dim() == 2: depth_map = depth_map.unsqueeze(0) # (H, W) -> (1, H, W) elif depth_map.dim() == 3: depth_map = depth_map.unsqueeze(1) # (B, H, W) -> (B, 1, H, W) elif depth_map.dim() != 4: raise ValueError(f"Unexpected depth shape: {original_shape}") invalid_mask = ~(torch.isfinite(depth_map) & (depth_map >= 0)) depth_map[invalid_mask] = 0.0 depth_map = depth_map / 1000.0 # mm -> meters max_depth_m = 11.2 depth_map = torch.clamp(depth_map, min=0.0, max=max_depth_m) depth_map = depth_map / max_depth_m # [0, 1] invalid_mask = ~torch.isfinite(depth_map) depth_map[invalid_mask] = 0.0 depth_map = depth_resize_transform(depth_map) # (..., 64, 128) depth_values = depth_map.squeeze(1) # [B, 64, 128] or [1, 64, 128] if len(original_shape) == 2: depth_values = depth_values.squeeze(0) # (64, 128) if is_numpy: depth_values = depth_values.numpy() return depth_values def quantize_depth_to_occupancy(depth_values, num_range_bins=64): """ Quantizes 2D depth values into 3D binary occupancy grid. Args: depth_values: Resized depth tensor Shape: [B, elevation, azimuth] Values: normalized to [0, 1] range num_range_bins: Number of range bins for quantization (default: 64) Returns: Binary 3D occupancy grid Shape: [B, elevation, azimuth, num_range_bins] Values: binary (0 or 1) indicating occupied bins """ B, elevation, azimuth = depth_values.shape # Quantize normalized depth [0, 1] directly to range bins [0, num_range_bins-1] # Each bin represents 1/num_range_bins of the normalized depth range bin_indices = torch.floor( depth_values / (1.0 / num_range_bins) ).long() # [B, elevation, azimuth] bin_indices = torch.clamp( bin_indices, 0, num_range_bins - 1 ) # Handle edge case where depth_values = 1.0 # Create binary 3D occupancy grid occupancy_grid = torch.zeros( B, elevation, azimuth, num_range_bins, dtype=torch.float32, device=depth_values.device, ) # [B, elevation, azimuth, num_range_bins] # Set occupied bins to 1 # Use advanced indexing to mark the appropriate range bin for each (elevation, azimuth) cell batch_idx = torch.arange(B, device=depth_values.device)[:, None, None].expand( B, elevation, azimuth ) elevation_idx = torch.arange(elevation, device=depth_values.device)[ None, :, None ].expand(B, elevation, azimuth) azimuth_idx = torch.arange(azimuth, device=depth_values.device)[ None, None, : ].expand(B, elevation, azimuth) occupancy_grid[batch_idx, elevation_idx, azimuth_idx, bin_indices] = 1.0 return occupancy_grid # [B, elevation, azimuth, num_range_bins] def dequantize_depth(occupancy_grid): """ Converts 3D binary occupancy grid back to 2D depth map. This is the inverse operation of quantize_depth_to_occupancy. Args: occupancy_grid: Binary 3D occupancy grid Shape: [B, 64, 128, 64] (batch, elevation, azimuth, range) Values: binary (0 or 1) or continuous (predicted probabilities) Returns: Reconstructed depth map Shape: [B, 1, 64, 128] (batch, channel, elevation, azimuth) Values: normalized to [0, 1] range """ num_range_bins = occupancy_grid.shape[3] # Find the range bin with maximum value for each (elevation, azimuth) cell # For binary: finds the occupied bin # For continuous: finds the most likely bin bin_indices = torch.argmax(occupancy_grid, dim=3) # [B, 64, 128] # Convert bin indices back to normalized depth values [0, 1] # Use bin center: (bin_idx + 0.5) / num_bins depth_values = (bin_indices.float() + 1) / num_range_bins # [B, 64, 128] # Add channel dimension: [B, 64, 128] -> [B, 1, 64, 128] depth_map = depth_values.unsqueeze(1) # [B, 1, 64, 128] return depth_map