Spaces:
Sleeping
Sleeping
File size: 2,747 Bytes
0cdb4e8 | 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 | """Floor/region segmentation (SAM) with a geometric fallback.
When a SAM checkpoint is configured (SAM_CHECKPOINT) and the ML stack is present,
this isolates the floor with point-prompted Segment Anything. Otherwise it
returns a robust geometric floor mask (region below the horizon, narrowing toward
the back wall). Either way the Shapely stage consumes a binary floor mask.
"""
from __future__ import annotations
from pathlib import Path
import cv2
import numpy as np
from ..config import settings
from .ml_runtime import ml_available
def floor_mask(image_path: Path, horizon_frac: float | None = None) -> np.ndarray:
"""Return a HxW uint8 mask (255 = floor) for the visible floor region."""
img = cv2.imread(str(image_path))
if img is None:
raise ValueError(f"Could not read image for segmentation: {image_path}")
h, w = img.shape[:2]
horizon_frac = settings.HORIZON_FRAC if horizon_frac is None else horizon_frac
if ml_available() and settings.SAM_CHECKPOINT:
m = _sam_floor_mask(img)
if m is not None:
return m
return geometric_floor_mask(h, w, horizon_frac)
def geometric_floor_mask(h: int, w: int, horizon_frac: float) -> np.ndarray:
"""Floor = trapezoid below the horizon, narrowing toward the back wall."""
horizon_y = int(h * horizon_frac)
poly = np.array(
[
[0, h - 1],
[w - 1, h - 1],
[int(w * 0.85), horizon_y],
[int(w * 0.15), horizon_y],
],
dtype=np.int32,
)
mask = np.zeros((h, w), dtype=np.uint8)
cv2.fillPoly(mask, [poly], 255)
return mask
def _sam_floor_mask(img: np.ndarray):
"""Point-prompted SAM floor isolation. Returns None on any problem."""
ckpt = Path(settings.SAM_CHECKPOINT)
if not ckpt.exists():
return None
try:
import torch
from segment_anything import SamPredictor, sam_model_registry
sam = sam_model_registry[settings.SAM_MODEL_TYPE](checkpoint=str(ckpt))
sam.to("cuda" if torch.cuda.is_available() else "cpu")
predictor = SamPredictor(sam)
rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
predictor.set_image(rgb)
h, w = img.shape[:2]
# Prompt with points along the bottom-center, where floor is most likely.
pts = np.array(
[[w // 2, int(h * 0.93)], [w // 3, int(h * 0.88)], [2 * w // 3, int(h * 0.88)]]
)
labels = np.ones(len(pts), dtype=np.int32)
masks, scores, _ = predictor.predict(
point_coords=pts, point_labels=labels, multimask_output=True
)
best = masks[int(np.argmax(scores))]
return (best.astype("uint8") * 255)
except Exception:
return None
|