Spaces:
Running on Zero
Running on Zero
File size: 3,648 Bytes
dc09dcc fb59dba dc09dcc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | # 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
|