Spaces:
Running on Zero
Running on Zero
| # Modified for SplatAnything: inference-only release. | |
| """ONNX sky segmentation (U^2-Net `skyseg.onnx`) enabled by --sky_seg. | |
| The inference entry point uses --sky_seg to enable both segmentation and cloud filtering. | |
| Per-image masks are marked with -1 in the input depth-mask tensor so the decoder can | |
| exclude sky pixels from scale estimation and exported Gaussians. | |
| Model contract (verified empirically against reference masks, IoU ~0.99): | |
| - input : NCHW float [1, 3, 320, 320] | |
| - output : the fused side map [1, 1, 320, 320] in [0, 1]; HIGH value (>0.5) = SKY. | |
| - preprocessing: RGB, resize to 320x320, /255, then ImageNet mean/std normalization. | |
| This is a per-image (perspective) segmenter. It intentionally does NOT implement the ERP-only | |
| GroundingDINO+ZIM pipeline (that one must see the full panorama); it is the pragmatic backend for | |
| ordinary perspective views where per-image inference is the only option. | |
| """ | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| _INPUT_SIZE = 320 | |
| _IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) | |
| _IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) | |
| class SkySegmenter: | |
| """Lazy ONNX sky segmenter. Call ``sky_mask(rgb)`` -> bool [H, W] (True = sky).""" | |
| def __init__(self, onnx_path: str, device: str = "cuda", prob_thresh: float = 0.5): | |
| self.onnx_path = str(onnx_path) | |
| self.prob_thresh = prob_thresh | |
| self.device = device | |
| self._sess = None | |
| self._iname = None | |
| self._oname = None | |
| if not Path(self.onnx_path).is_file(): | |
| raise FileNotFoundError(f"skyseg ONNX model not found: {self.onnx_path}") | |
| def _ensure_session(self): | |
| if self._sess is not None: | |
| return | |
| import onnxruntime as ort # local import: optional dependency | |
| providers = (["CUDAExecutionProvider", "CPUExecutionProvider"] | |
| if self.device.startswith("cuda") else ["CPUExecutionProvider"]) | |
| options = ort.SessionOptions() | |
| # Keep per-image CPU segmentation bounded on shared Space CPU cores. | |
| options.intra_op_num_threads = 2 | |
| options.inter_op_num_threads = 1 | |
| self._sess = ort.InferenceSession(self.onnx_path, sess_options=options, providers=providers) | |
| self._iname = self._sess.get_inputs()[0].name | |
| self._oname = self._sess.get_outputs()[0].name # fused map (first output) | |
| print(f"[SkySeg] loaded {self.onnx_path} on {self._sess.get_providers()[0]}") | |
| def sky_mask(self, rgb: np.ndarray) -> np.ndarray: | |
| """rgb: [H, W, 3] uint8 or float in [0, 1]/[0, 255]. Returns bool [H, W], True = sky. | |
| The 320x320 probability map is resized back to (H, W) with NEAREST so the binary mask | |
| stays crisp and aligns 1:1 with the per-pixel Gaussian grid. | |
| """ | |
| self._ensure_session() | |
| rgb = np.asarray(rgb) | |
| if rgb.dtype != np.uint8: | |
| # accept [0,1] floats; scale to [0,255] for the /255 preprocessing below | |
| rgb = (rgb * 255.0).clip(0, 255) if rgb.max() <= 1.0 + 1e-6 else rgb | |
| H, W = rgb.shape[:2] | |
| im = cv2.resize(rgb.astype(np.float32), (_INPUT_SIZE, _INPUT_SIZE), | |
| interpolation=cv2.INTER_AREA) | |
| im = (im / 255.0 - _IMAGENET_MEAN) / _IMAGENET_STD | |
| x = im.transpose(2, 0, 1)[None].astype(np.float32) # [1, 3, 320, 320] | |
| prob = self._sess.run([self._oname], {self._iname: x})[0][0, 0] # [320, 320], high = sky | |
| sky_small = (prob > self.prob_thresh).astype(np.uint8) | |
| sky = cv2.resize(sky_small, (W, H), interpolation=cv2.INTER_NEAREST) > 0 | |
| return sky | |