Spaces:
Sleeping
Sleeping
| """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 | |