# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. import logging import copy import os import cv2 import matplotlib import numpy as np import requests import trimesh from scipy.spatial import cKDTree from scipy.spatial.transform import Rotation logger = logging.getLogger(__name__) def _srgb_to_linear(colors: np.ndarray) -> np.ndarray: colors = np.clip(colors, 0.0, 1.0) threshold = 0.04045 below = colors <= threshold linear = np.empty_like(colors, dtype=np.float64) linear[below] = colors[below] / 12.92 linear[~below] = ((colors[~below] + 0.055) / 1.055) ** 2.4 return linear def _linear_to_srgb(colors: np.ndarray) -> np.ndarray: colors = np.clip(colors, 0.0, 1.0) threshold = 0.0031308 srgb = np.empty_like(colors, dtype=np.float64) below = colors <= threshold srgb[below] = colors[below] * 12.92 srgb[~below] = 1.055 * np.power(colors[~below], 1 / 2.4) - 0.055 return np.clip(np.round(srgb * 255.0), 0, 255).astype(np.uint8) def voxel_reduce( points_f32: np.ndarray, colors_u8: np.ndarray, conf_f32: np.ndarray | None = None, voxel_size: float = 0.02, origin: np.ndarray | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: points = np.asarray(points_f32, dtype=np.float32) colors = np.asarray(colors_u8, dtype=np.uint8) if points.size == 0: return ( points.reshape(-1, 3).astype(np.float32), colors.reshape(-1, 3).astype(np.uint8), np.zeros((points.shape[0],), dtype=np.float32), ) if voxel_size is None or voxel_size <= 0: weights = ( np.asarray(conf_f32, dtype=np.float32).reshape(-1) if conf_f32 is not None else np.ones(points.shape[0], dtype=np.float32) ) return points.astype(np.float32), colors.astype(np.uint8), weights weights = ( np.asarray(conf_f32, dtype=np.float32).reshape(-1) if conf_f32 is not None else np.ones(points.shape[0], dtype=np.float32) ) if weights.shape[0] != points.shape[0]: raise ValueError("conf_f32 must match the shape of points.") base = ( np.asarray(origin, dtype=np.float32) if origin is not None else points.min(axis=0).astype(np.float32) ) voxel_indices = np.floor((points - base) / voxel_size).astype(np.int64) voxel_keys, inverse_indices, counts = np.unique( voxel_indices, axis=0, return_inverse=True, return_counts=True ) reduced_count = voxel_keys.shape[0] accum_weights = np.bincount(inverse_indices, weights=weights, minlength=reduced_count) accum_weights = np.where(accum_weights <= 0, 1e-6, accum_weights) reduced_points = np.zeros((reduced_count, 3), dtype=np.float64) for dim in range(3): reduced_points[:, dim] = np.bincount( inverse_indices, weights=weights * points[:, dim], minlength=reduced_count, ) reduced_points /= accum_weights[:, None] colors_linear = _srgb_to_linear(colors.astype(np.float32) / 255.0) reduced_colors_linear = np.zeros((reduced_count, 3), dtype=np.float64) for dim in range(3): reduced_colors_linear[:, dim] = np.bincount( inverse_indices, weights=weights * colors_linear[:, dim], minlength=reduced_count, ) reduced_colors_linear /= accum_weights[:, None] reduced_colors = _linear_to_srgb(reduced_colors_linear) support = ( accum_weights.astype(np.float32) if conf_f32 is not None else counts.astype(np.float32) ) return reduced_points.astype(np.float32), reduced_colors.astype(np.uint8), support def _filter_by_support( points: np.ndarray, colors: np.ndarray, support: np.ndarray, min_support: float | None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: if ( support is None or support.size == 0 or min_support is None or min_support <= 0 ): return points, colors, support mask = support >= float(min_support) if not np.any(mask): return points, colors, support return points[mask], colors[mask], support[mask] def _log_point_count(stage: str, before: int, after: int) -> None: if logger.isEnabledFor(logging.INFO): logger.info("Point cloud %s: %d -> %d", stage, before, after) def o3d_outlier_filter( points_f32: np.ndarray, colors_u8: np.ndarray, *, voxel_size: float = 0.02, radius_mult: float = 3.0, nb_points: int = 16, nb_neighbors: int = 48, std_ratio: float = 1.5, ) -> tuple[np.ndarray, np.ndarray]: points = np.asarray(points_f32, dtype=np.float32) colors = np.asarray(colors_u8, dtype=np.uint8) if points.size == 0: return points.reshape(-1, 3), colors.reshape(-1, 3) try: import open3d as o3d # type: ignore except ImportError: logger.warning("Open3D not available; skipping outlier filtering.") return points.astype(np.float32), colors.astype(np.uint8) pcd = o3d.geometry.PointCloud() pcd.points = o3d.utility.Vector3dVector(points.astype(np.float64)) pcd.colors = o3d.utility.Vector3dVector(colors.astype(np.float32) / 255.0) effective_voxel = float(voxel_size) if voxel_size and voxel_size > 0 else 0.02 radius = max(float(radius_mult) * effective_voxel, 1e-4) if nb_points > 0: pcd, _ = pcd.remove_radius_outlier(nb_points=int(nb_points), radius=radius) if len(pcd.points) == 0: return np.empty((0, 3), dtype=np.float32), np.empty((0, 3), dtype=np.uint8) if nb_neighbors > 0: pcd, _ = pcd.remove_statistical_outlier( nb_neighbors=int(nb_neighbors), std_ratio=float(std_ratio), ) if len(pcd.points) == 0: return np.empty((0, 3), dtype=np.float32), np.empty((0, 3), dtype=np.uint8) filtered_points = np.asarray(pcd.points, dtype=np.float32) filtered_colors = np.asarray(pcd.colors, dtype=np.float32) filtered_colors = np.clip(np.round(filtered_colors * 255.0), 0, 255).astype(np.uint8) return filtered_points, filtered_colors def density_filter_points( points_f32: np.ndarray, colors_u8: np.ndarray, *, radius: float, min_neighbors: int, ) -> tuple[np.ndarray, np.ndarray]: points = np.asarray(points_f32, dtype=np.float32) colors = np.asarray(colors_u8, dtype=np.uint8) if points.size == 0: return points.reshape(-1, 3), colors.reshape(-1, 3) radius = max(float(radius), 1e-4) min_neighbors = max(int(min_neighbors), 1) tree = cKDTree(points) neighbor_lists = tree.query_ball_point(points, radius) mask = np.fromiter((len(nlist) >= min_neighbors for nlist in neighbor_lists), dtype=bool, count=len(neighbor_lists)) return points[mask], colors[mask] def reinflate_voxels( points: np.ndarray, colors: np.ndarray, support: np.ndarray | None, *, voxel_size: float, support_scale: float = 0.5, min_samples: int = 1, max_samples: int | None = 12, jitter_mode: str = "cube", jitter_sigma: float = 0.35, seed: int | None = None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ Expand each voxel centroid into a jittered micro-cluster to recover splat coverage. Returns expanded points, colors, and per-sample support values whose sum matches the original. """ if ( points.size == 0 or support is None or voxel_size is None or voxel_size <= 0 ): return points, colors, np.asarray(support) if support is not None else np.array([], dtype=np.float32) support = np.asarray(support, dtype=np.float32) if support.shape[0] != points.shape[0]: raise ValueError("Support array must align with points for reinflation.") voxel_size = float(voxel_size) positive_mask = support > 0 if not np.any(positive_mask): return points, colors, support counts = np.zeros_like(support, dtype=np.int32) scaled = np.round(support_scale * support[positive_mask]).astype(np.int32) if min_samples is not None: scaled = np.maximum(scaled, int(max(0, min_samples))) if max_samples is not None: scaled = np.minimum(scaled, int(max_samples)) scaled = np.maximum(scaled, 0) counts[positive_mask] = scaled total_samples = int(counts.sum()) if total_samples == 0: fallback = max(1, int(min_samples or 1)) if max_samples is not None: fallback = min(fallback, int(max_samples)) counts[positive_mask] = fallback total_samples = int(counts.sum()) rng = np.random.default_rng(seed) repeated_indices = np.repeat(np.arange(points.shape[0]), counts) expanded_points = points[repeated_indices].astype(np.float32, copy=True) expanded_colors = colors[repeated_indices].astype(np.uint8, copy=True) offsets = np.zeros_like(expanded_points, dtype=np.float32) jitter_sigma = float(max(jitter_sigma, 0.0)) if jitter_mode not in {"cube", "gaussian"}: raise ValueError("jitter_mode must be 'cube' or 'gaussian'") if jitter_sigma > 0 and total_samples > 0: if jitter_mode == "cube": span = 0.5 * voxel_size * jitter_sigma offsets = rng.uniform(-span, span, size=expanded_points.shape).astype(np.float32) else: sigma = voxel_size * jitter_sigma offsets = rng.normal(0.0, sigma, size=expanded_points.shape).astype(np.float32) max_span = 0.5 * voxel_size np.clip(offsets, -max_span, max_span, out=offsets) cumulative = np.cumsum(counts) starts = cumulative - counts offsets[starts] = 0.0 expanded_points += offsets normalized_support = np.zeros_like(support, dtype=np.float32) nonzero_counts = counts > 0 normalized_support[nonzero_counts] = support[nonzero_counts] / counts[nonzero_counts] expanded_support = np.repeat(normalized_support, counts) return expanded_points, expanded_colors, expanded_support def predictions_to_glb( predictions, conf_thres=50.0, filter_by_frames="all", mask_black_bg=False, mask_white_bg=False, show_cam=True, mask_sky=False, target_dir=None, prediction_mode="Predicted Pointmap", extra_cameras=None, extra_camera_color=(255, 0, 0), voxel_size: float | None = 0.01, voxel_after_conf: bool = True, min_voxel_support: float | None = 3, o3d_denoise: bool = True, o3d_params: dict | None = None, density_filter: bool = False, density_params: dict | None = None, reinflate_enabled: bool = True, reinflate_support_scale: float = 1.5, reinflate_min_samples: int = 3, reinflate_max_samples: int | None = 8, reinflate_jitter_mode: str = "cube", reinflate_jitter_sigma: float = 0.35, reinflate_seed: int | None = None, ceiling_percentile: float | None = None, ceiling_margin: float = 0.05, ceiling_z_max: float | None = None, ) -> trimesh.Scene: """ Converts predictions to a 3D scene represented as a GLB file. Args: predictions (dict): Dictionary containing model predictions with keys: - world_points: 3D point coordinates (S, H, W, 3) - world_points_conf: Confidence scores (S, H, W) - images: Input images (S, H, W, 3) - extrinsic: Camera extrinsic matrices (S, 3, 4) conf_thres (float): Percentage of low-confidence points to filter out (default: 50.0) filter_by_frames (str): Frame filter specification (default: "all") mask_black_bg (bool): Mask out black background pixels (default: False) mask_white_bg (bool): Mask out white background pixels (default: False) show_cam (bool): Include camera visualization (default: True) mask_sky (bool): Apply sky segmentation mask (default: False) target_dir (str): Output directory for intermediate files (default: None) prediction_mode (str): Prediction mode selector (default: "Predicted Pointmap") extra_cameras (Optional[List[np.ndarray]]): Additional camera extrinsics (3x4 or 4x4) to visualize even when show_cam=False. Useful for highlighting localized poses. extra_camera_color (tuple or list[tuple]): RGB color(s) for extra cameras. voxel_size (Optional[float]): Size of voxel grid cells (>0 enables reduction). voxel_after_conf (bool): Apply voxel reduction after confidence/background filtering. min_voxel_support (Optional[float]): Minimum aggregated support (confidence/count) per voxel. o3d_denoise (bool): Enable Open3D outlier filtering. o3d_params (Optional[dict]): Overrides for Open3D filtering parameters. density_filter (bool): Apply KD-tree based density filtering. density_params (Optional[dict]): Overrides for density filter parameters. reinflate_enabled (bool): Re-expand voxels into jittered micro-clusters. reinflate_support_scale (float): Multiplier converting support into sample count. reinflate_min_samples (int): Minimum samples emitted per voxel with positive support. reinflate_max_samples (Optional[int]): Maximum samples emitted per voxel. reinflate_jitter_mode (str): "cube" (uniform jitter) or "gaussian". reinflate_jitter_sigma (float): Jitter strength as a fraction of voxel size. reinflate_seed (Optional[int]): RNG seed for deterministic reinflation. ceiling_percentile (Optional[float]): Remove points above this Z percentile (0-100). ceiling_margin (float): Margin subtracted from percentile cutoff (meters). ceiling_z_max (Optional[float]): Remove points with Z >= this absolute height (meters). Returns: trimesh.Scene: Processed 3D scene containing point cloud and cameras Raises: ValueError: If input predictions structure is invalid """ if not isinstance(predictions, dict): raise ValueError("predictions must be a dictionary") if conf_thres is None: conf_thres = 10.0 print("Building GLB scene") selected_frame_idx = None if filter_by_frames != "all" and filter_by_frames != "All": try: # Extract the index part before the colon selected_frame_idx = int(filter_by_frames.split(":")[0]) except (ValueError, IndexError): pass if "Pointmap" in prediction_mode: print("Using Pointmap Branch") if "world_points" in predictions: pred_world_points = predictions["world_points"] # No batch dimension to remove pred_world_points_conf = predictions.get("world_points_conf", np.ones_like(pred_world_points[..., 0])) else: print("Warning: world_points not found in predictions, falling back to depth-based points") pred_world_points = predictions["world_points_from_depth"] pred_world_points_conf = predictions.get("depth_conf", np.ones_like(pred_world_points[..., 0])) else: print("Using Depthmap and Camera Branch") pred_world_points = predictions["world_points_from_depth"] pred_world_points_conf = predictions.get("depth_conf", np.ones_like(pred_world_points[..., 0])) # Get images from predictions images = predictions["images"] # Use extrinsic matrices instead of pred_extrinsic_list camera_matrices = predictions["extrinsic"] if mask_sky: if target_dir is not None: import onnxruntime skyseg_session = None target_dir_images = target_dir + "/images" image_list = sorted(os.listdir(target_dir_images)) sky_mask_list = [] # Get the shape of pred_world_points_conf to match S, H, W = ( pred_world_points_conf.shape if hasattr(pred_world_points_conf, "shape") else (len(images), images.shape[1], images.shape[2]) ) # Download skyseg.onnx if it doesn't exist if not os.path.exists("skyseg.onnx"): print("Downloading skyseg.onnx...") download_file_from_url( "https://huggingface.co/JianyuanWang/skyseg/resolve/main/skyseg.onnx", "skyseg.onnx" ) for i, image_name in enumerate(image_list): image_filepath = os.path.join(target_dir_images, image_name) mask_filepath = os.path.join(target_dir, "sky_masks", image_name) # Check if mask already exists if os.path.exists(mask_filepath): # Load existing mask sky_mask = cv2.imread(mask_filepath, cv2.IMREAD_GRAYSCALE) else: # Generate new mask if skyseg_session is None: skyseg_session = onnxruntime.InferenceSession("skyseg.onnx") sky_mask = segment_sky(image_filepath, skyseg_session, mask_filepath) # Resize mask to match H×W if needed if sky_mask.shape[0] != H or sky_mask.shape[1] != W: sky_mask = cv2.resize(sky_mask, (W, H)) sky_mask_list.append(sky_mask) # Convert list to numpy array with shape S×H×W sky_mask_array = np.array(sky_mask_list) # Apply sky mask to confidence scores sky_mask_binary = (sky_mask_array > 0.1).astype(np.float32) pred_world_points_conf = pred_world_points_conf * sky_mask_binary if selected_frame_idx is not None: pred_world_points = pred_world_points[selected_frame_idx][None] pred_world_points_conf = pred_world_points_conf[selected_frame_idx][None] images = images[selected_frame_idx][None] camera_matrices = camera_matrices[selected_frame_idx][None] vertices_3d = pred_world_points.reshape(-1, 3) # Handle different image formats - check if images need transposing if images.ndim == 4 and images.shape[1] == 3: # NCHW format colors_rgb = np.transpose(images, (0, 2, 3, 1)) else: # Assume already in NHWC format colors_rgb = images colors_rgb = (colors_rgb.reshape(-1, 3) * 255).astype(np.uint8) conf = pred_world_points_conf.reshape(-1).astype(np.float32) effective_voxel_size = float(voxel_size) if voxel_size is not None else None if effective_voxel_size is not None and effective_voxel_size <= 0: effective_voxel_size = None if effective_voxel_size is not None and not voxel_after_conf: before_count = vertices_3d.shape[0] vertices_3d, colors_rgb, conf = voxel_reduce( vertices_3d, colors_rgb, conf, voxel_size=effective_voxel_size, ) vertices_3d, colors_rgb, conf = _filter_by_support( vertices_3d, colors_rgb, conf, min_voxel_support, ) after_reduce = vertices_3d.shape[0] _log_point_count("voxel_reduce_pre_conf", before_count, after_reduce) if reinflate_enabled and after_reduce: vertices_3d, colors_rgb, conf = reinflate_voxels( vertices_3d, colors_rgb, conf, voxel_size=effective_voxel_size, support_scale=reinflate_support_scale, min_samples=reinflate_min_samples, max_samples=reinflate_max_samples, jitter_mode=reinflate_jitter_mode, jitter_sigma=reinflate_jitter_sigma, seed=reinflate_seed, ) _log_point_count("voxel_reinflate_pre_conf", after_reduce, vertices_3d.shape[0]) # Convert percentage threshold to actual confidence value if conf_thres == 0.0: conf_threshold = 0.0 else: conf_threshold = np.percentile(conf, conf_thres) conf_mask = (conf >= conf_threshold) & (conf > 1e-5) if mask_black_bg: black_bg_mask = colors_rgb.sum(axis=1) >= 16 conf_mask = conf_mask & black_bg_mask if mask_white_bg: # Filter out white background pixels (RGB values close to white) # Consider pixels white if all RGB values are above 240 white_bg_mask = ~((colors_rgb[:, 0] > 240) & (colors_rgb[:, 1] > 240) & (colors_rgb[:, 2] > 240)) conf_mask = conf_mask & white_bg_mask vertices_3d = vertices_3d[conf_mask] colors_rgb = colors_rgb[conf_mask] conf_used = conf[conf_mask] if ceiling_percentile is not None and vertices_3d.size: try: percentile_value = float(ceiling_percentile) except (TypeError, ValueError): percentile_value = None if percentile_value is not None and 0.0 < percentile_value < 100.0: cutoff = float(np.percentile(vertices_3d[:, 2], percentile_value)) margin = float(max(0.0, ceiling_margin)) threshold = cutoff - margin keep_mask = vertices_3d[:, 2] < threshold if not np.any(keep_mask): keep_mask = vertices_3d[:, 2] <= cutoff if np.any(keep_mask) and np.count_nonzero(keep_mask) < vertices_3d.shape[0]: vertices_3d = vertices_3d[keep_mask] colors_rgb = colors_rgb[keep_mask] conf_used = conf_used[keep_mask] if ceiling_z_max is not None and vertices_3d.size: try: z_limit = float(ceiling_z_max) except (TypeError, ValueError): z_limit = None if z_limit is not None: keep_mask = vertices_3d[:, 2] < z_limit if not np.any(keep_mask): keep_mask = vertices_3d[:, 2] <= z_limit if np.any(keep_mask) and np.count_nonzero(keep_mask) < vertices_3d.shape[0]: vertices_3d = vertices_3d[keep_mask] colors_rgb = colors_rgb[keep_mask] conf_used = conf_used[keep_mask] if effective_voxel_size is not None and voxel_after_conf and vertices_3d.size: before_count = vertices_3d.shape[0] vertices_3d, colors_rgb, conf_used = voxel_reduce( vertices_3d, colors_rgb, conf_used, voxel_size=effective_voxel_size, ) vertices_3d, colors_rgb, conf_used = _filter_by_support( vertices_3d, colors_rgb, conf_used, min_voxel_support, ) after_reduce = vertices_3d.shape[0] _log_point_count("voxel_reduce_post_conf", before_count, after_reduce) if reinflate_enabled and after_reduce: vertices_3d, colors_rgb, conf_used = reinflate_voxels( vertices_3d, colors_rgb, conf_used, voxel_size=effective_voxel_size, support_scale=reinflate_support_scale, min_samples=reinflate_min_samples, max_samples=reinflate_max_samples, jitter_mode=reinflate_jitter_mode, jitter_sigma=reinflate_jitter_sigma, seed=reinflate_seed, ) _log_point_count("voxel_reinflate_post_conf", after_reduce, vertices_3d.shape[0]) if o3d_denoise and vertices_3d.size: before_count = vertices_3d.shape[0] params = { "voxel_size": effective_voxel_size or 0.02, "radius_mult": 3.0, "nb_points": 16, "nb_neighbors": 48, "std_ratio": 1.5, } if o3d_params: params.update(o3d_params) vertices_3d, colors_rgb = o3d_outlier_filter(vertices_3d, colors_rgb, **params) _log_point_count("o3d_denoise", before_count, vertices_3d.shape[0]) if density_filter and vertices_3d.size: before_count = vertices_3d.shape[0] params = { "radius": (effective_voxel_size or 0.02) * 2.5, "min_neighbors": 6, } if density_params: params.update(density_params) vertices_3d, colors_rgb = density_filter_points(vertices_3d, colors_rgb, **params) _log_point_count("density_filter", before_count, vertices_3d.shape[0]) if vertices_3d is None or np.asarray(vertices_3d).size == 0: vertices_3d = np.array([[1, 0, 0]]) colors_rgb = np.array([[255, 255, 255]]) scene_scale = 1 else: # Calculate the 5th and 95th percentiles along each axis lower_percentile = np.percentile(vertices_3d, 5, axis=0) upper_percentile = np.percentile(vertices_3d, 95, axis=0) # Calculate the diagonal length of the percentile bounding box scene_scale = np.linalg.norm(upper_percentile - lower_percentile) colormap = matplotlib.colormaps.get_cmap("gist_rainbow") # Initialize a 3D scene scene_3d = trimesh.Scene() # Add point cloud data to the scene point_cloud_data = trimesh.PointCloud(vertices=vertices_3d, colors=colors_rgb) scene_3d.add_geometry(point_cloud_data) # Prepare 4x4 matrices for camera extrinsics num_cameras = len(camera_matrices) extrinsics_matrices = np.zeros((num_cameras, 4, 4)) extrinsics_matrices[:, :3, :4] = camera_matrices extrinsics_matrices[:, 3, 3] = 1 extra_cameras = [] if extra_cameras is None else list(extra_cameras) if isinstance(extra_camera_color, tuple) and len(extra_cameras) > 1: extra_colors = [extra_camera_color for _ in extra_cameras] elif isinstance(extra_camera_color, (list, tuple)) and len(extra_cameras) == len(extra_camera_color): extra_colors = list(extra_camera_color) else: extra_colors = [(255, 0, 0) for _ in extra_cameras] if show_cam: # Add camera models to the scene for i in range(num_cameras): world_to_camera = extrinsics_matrices[i] camera_to_world = np.linalg.inv(world_to_camera) rgba_color = colormap(i / num_cameras) current_color = tuple(int(255 * x) for x in rgba_color[:3]) integrate_camera_into_scene(scene_3d, camera_to_world, current_color, scene_scale) for idx, extra in enumerate(extra_cameras): extra = np.asarray(extra) if extra.shape == (3, 4): world_to_camera = np.eye(4) world_to_camera[:3, :4] = extra elif extra.shape == (4, 4): world_to_camera = extra else: raise ValueError("Extra camera extrinsic must have shape (3,4) or (4,4)") camera_to_world = np.linalg.inv(world_to_camera) integrate_camera_into_scene( scene_3d, camera_to_world, extra_colors[idx] if idx < len(extra_colors) else (255, 0, 0), scene_scale, ) # Align scene to the observation of the first camera scene_3d = apply_scene_alignment(scene_3d, extrinsics_matrices) print("GLB Scene built") return scene_3d def integrate_camera_into_scene(scene: trimesh.Scene, transform: np.ndarray, face_colors: tuple, scene_scale: float): """ Integrates a fake camera mesh into the 3D scene. Args: scene (trimesh.Scene): The 3D scene to add the camera model. transform (np.ndarray): Transformation matrix for camera positioning. face_colors (tuple): Color of the camera face. scene_scale (float): Scale of the scene. """ cam_width = scene_scale * 0.05 cam_height = scene_scale * 0.1 # Create cone shape for camera rot_45_degree = np.eye(4) rot_45_degree[:3, :3] = Rotation.from_euler("z", 45, degrees=True).as_matrix() rot_45_degree[2, 3] = -cam_height opengl_transform = get_opengl_conversion_matrix() # Combine transformations complete_transform = transform @ opengl_transform @ rot_45_degree camera_cone_shape = trimesh.creation.cone(cam_width, cam_height, sections=4) # Generate mesh for the camera slight_rotation = np.eye(4) slight_rotation[:3, :3] = Rotation.from_euler("z", 2, degrees=True).as_matrix() vertices_combined = np.concatenate( [ camera_cone_shape.vertices, 0.95 * camera_cone_shape.vertices, transform_points(slight_rotation, camera_cone_shape.vertices), ] ) vertices_transformed = transform_points(complete_transform, vertices_combined) mesh_faces = compute_camera_faces(camera_cone_shape) # Add the camera mesh to the scene camera_mesh = trimesh.Trimesh(vertices=vertices_transformed, faces=mesh_faces) camera_mesh.visual.face_colors[:, :3] = face_colors scene.add_geometry(camera_mesh) def apply_scene_alignment(scene_3d: trimesh.Scene, extrinsics_matrices: np.ndarray) -> trimesh.Scene: """ Aligns the 3D scene based on the extrinsics of the first camera. Args: scene_3d (trimesh.Scene): The 3D scene to be aligned. extrinsics_matrices (np.ndarray): Camera extrinsic matrices. Returns: trimesh.Scene: Aligned 3D scene. """ # Set transformations for scene alignment opengl_conversion_matrix = get_opengl_conversion_matrix() # Rotation matrix for alignment (180 degrees around the y-axis) align_rotation = np.eye(4) align_rotation[:3, :3] = Rotation.from_euler("y", 180, degrees=True).as_matrix() # Apply transformation initial_transformation = np.linalg.inv(extrinsics_matrices[0]) @ opengl_conversion_matrix @ align_rotation scene_3d.apply_transform(initial_transformation) return scene_3d def get_opengl_conversion_matrix() -> np.ndarray: """ Constructs and returns the OpenGL conversion matrix. Returns: numpy.ndarray: A 4x4 OpenGL conversion matrix. """ # Create an identity matrix matrix = np.identity(4) # Flip the y and z axes matrix[1, 1] = -1 matrix[2, 2] = -1 return matrix def transform_points(transformation: np.ndarray, points: np.ndarray, dim: int = None) -> np.ndarray: """ Applies a 4x4 transformation to a set of points. Args: transformation (np.ndarray): Transformation matrix. points (np.ndarray): Points to be transformed. dim (int, optional): Dimension for reshaping the result. Returns: np.ndarray: Transformed points. """ points = np.asarray(points) initial_shape = points.shape[:-1] dim = dim or points.shape[-1] # Apply transformation transformation = transformation.swapaxes(-1, -2) # Transpose the transformation matrix points = points @ transformation[..., :-1, :] + transformation[..., -1:, :] # Reshape the result result = points[..., :dim].reshape(*initial_shape, dim) return result def compute_camera_faces(cone_shape: trimesh.Trimesh) -> np.ndarray: """ Computes the faces for the camera mesh. Args: cone_shape (trimesh.Trimesh): The shape of the camera cone. Returns: np.ndarray: Array of faces for the camera mesh. """ # Create pseudo cameras faces_list = [] num_vertices_cone = len(cone_shape.vertices) for face in cone_shape.faces: if 0 in face: continue v1, v2, v3 = face v1_offset, v2_offset, v3_offset = face + num_vertices_cone v1_offset_2, v2_offset_2, v3_offset_2 = face + 2 * num_vertices_cone faces_list.extend( [ (v1, v2, v2_offset), (v1, v1_offset, v3), (v3_offset, v2, v3), (v1, v2, v2_offset_2), (v1, v1_offset_2, v3), (v3_offset_2, v2, v3), ] ) faces_list += [(v3, v2, v1) for v1, v2, v3 in faces_list] return np.array(faces_list) def segment_sky(image_path, onnx_session, mask_filename=None): """ Segments sky from an image using an ONNX model. Thanks for the great model provided by https://github.com/xiongzhu666/Sky-Segmentation-and-Post-processing Args: image_path: Path to input image onnx_session: ONNX runtime session with loaded model mask_filename: Path to save the output mask Returns: np.ndarray: Binary mask where 255 indicates non-sky regions """ assert mask_filename is not None image = cv2.imread(image_path) result_map = run_skyseg(onnx_session, [320, 320], image) # resize the result_map to the original image size result_map_original = cv2.resize(result_map, (image.shape[1], image.shape[0])) # Fix: Invert the mask so that 255 = non-sky, 0 = sky # The model outputs low values for sky, high values for non-sky output_mask = np.zeros_like(result_map_original) output_mask[result_map_original < 32] = 255 # Use threshold of 32 os.makedirs(os.path.dirname(mask_filename), exist_ok=True) cv2.imwrite(mask_filename, output_mask) return output_mask def run_skyseg(onnx_session, input_size, image): """ Runs sky segmentation inference using ONNX model. Args: onnx_session: ONNX runtime session input_size: Target size for model input (width, height) image: Input image in BGR format Returns: np.ndarray: Segmentation mask """ # Pre process:Resize, BGR->RGB, Transpose, PyTorch standardization, float32 cast temp_image = copy.deepcopy(image) resize_image = cv2.resize(temp_image, dsize=(input_size[0], input_size[1])) x = cv2.cvtColor(resize_image, cv2.COLOR_BGR2RGB) x = np.array(x, dtype=np.float32) mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] x = (x / 255 - mean) / std x = x.transpose(2, 0, 1) x = x.reshape(-1, 3, input_size[0], input_size[1]).astype("float32") # Inference input_name = onnx_session.get_inputs()[0].name output_name = onnx_session.get_outputs()[0].name onnx_result = onnx_session.run([output_name], {input_name: x}) # Post process onnx_result = np.array(onnx_result).squeeze() min_value = np.min(onnx_result) max_value = np.max(onnx_result) onnx_result = (onnx_result - min_value) / (max_value - min_value) onnx_result *= 255 onnx_result = onnx_result.astype("uint8") return onnx_result def download_file_from_url(url, filename): """Downloads a file from a Hugging Face model repo, handling redirects.""" try: # Get the redirect URL response = requests.get(url, allow_redirects=False) response.raise_for_status() # Raise HTTPError for bad requests (4xx or 5xx) if response.status_code == 302: # Expecting a redirect redirect_url = response.headers["Location"] response = requests.get(redirect_url, stream=True) response.raise_for_status() else: print(f"Unexpected status code: {response.status_code}") return with open(filename, "wb") as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) print(f"Downloaded {filename} successfully.") except requests.exceptions.RequestException as e: print(f"Error downloading file: {e}")