Spaces:
Sleeping
Sleeping
RiceZero
fix: enhance FastSAM model loading logic to prioritize local path and fallback to cache
7b9fc68 | """Object segmentation using FastSAM.""" | |
| from __future__ import annotations | |
| import gc | |
| import logging | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| from src.models.segment import Segment | |
| logger = logging.getLogger(__name__) | |
| class SegmentationService: | |
| """Segments objects in images using FastSAM (CPU-friendly).""" | |
| def __init__( | |
| self, | |
| model_path: str, | |
| conf: float = 0.4, | |
| iou: float = 0.9, | |
| imgsz: int = 1024, | |
| ) -> None: | |
| self._model_path = model_path | |
| self._conf = conf | |
| self._iou = iou | |
| self._imgsz = imgsz | |
| self._model = None # lazy load | |
| def _ensure_model(self) -> None: | |
| """Load the model on first use. | |
| Tries to load from local path first, falls back to Ultralytics cache. | |
| """ | |
| if self._model is not None: | |
| return | |
| from ultralytics import FastSAM | |
| path = Path(self._model_path) | |
| # Try local path first | |
| if path.exists() and path.stat().st_size > 1000: # More than 1KB | |
| logger.info("Loading FastSAM model from %s ...", path) | |
| self._model = FastSAM(str(path)) | |
| logger.info("FastSAM model loaded successfully from local path.") | |
| else: | |
| # Fall back to Ultralytics automatic download/cache | |
| logger.info("Local model not found, using Ultralytics cache for FastSAM-x...") | |
| logger.info("This will download the model (~1.3GB) on first use...") | |
| self._model = FastSAM('FastSAM-x.pt') # Ultralytics handles download | |
| logger.info("FastSAM model loaded successfully from Ultralytics cache.") | |
| def unload_model(self) -> None: | |
| """Unload model from memory to free up space. | |
| Call this method when you need to free memory in a 4GB environment. | |
| The model will be reloaded automatically on next segment() call. | |
| """ | |
| if self._model is not None: | |
| logger.info("Unloading FastSAM model to free memory...") | |
| self._model = None | |
| gc.collect() | |
| logger.info("FastSAM model unloaded.") | |
| def segment( | |
| self, | |
| image: np.ndarray, | |
| min_area_ratio: float = 0.002, | |
| max_area_ratio: float = 0.15, | |
| ) -> list[Segment]: | |
| """Run FastSAM segmentation on a BGR image. | |
| Args: | |
| image: BGR image (H, W, 3) uint8. | |
| min_area_ratio: Minimum segment area as fraction of image area. | |
| max_area_ratio: Maximum segment area as fraction of image area. | |
| Returns: | |
| List of Segment objects, sorted by area (descending). | |
| """ | |
| self._ensure_model() | |
| h, w = image.shape[:2] | |
| total_pixels = h * w | |
| min_area = int(total_pixels * min_area_ratio) | |
| max_area = int(total_pixels * max_area_ratio) | |
| # FastSAM expects RGB or file path; convert BGR -> RGB | |
| rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| results = self._model( | |
| rgb, | |
| device="cpu", | |
| retina_masks=True, | |
| imgsz=self._imgsz, | |
| conf=self._conf, | |
| iou=self._iou, | |
| verbose=False, | |
| ) | |
| segments = self._extract_segments(results, h, w, min_area, max_area) | |
| # Memory cleanup for 4GB environment | |
| # Clear temporary data after segmentation | |
| del rgb | |
| del results | |
| gc.collect() | |
| logger.info( | |
| "Segmentation complete: %d segments (filtered from raw results)", | |
| len(segments), | |
| ) | |
| return segments | |
| def _extract_segments( | |
| self, | |
| results, | |
| img_h: int, | |
| img_w: int, | |
| min_area: int, | |
| max_area: int, | |
| ) -> list[Segment]: | |
| """Extract Segment objects from FastSAM results with improved filtering.""" | |
| segments: list[Segment] = [] | |
| if not results or results[0].masks is None: | |
| logger.warning("No masks returned by FastSAM.") | |
| return segments | |
| masks_data = results[0].masks.data.cpu().numpy() | |
| boxes = results[0].boxes | |
| for i in range(len(masks_data)): | |
| mask = masks_data[i] | |
| # Resize mask to original image dimensions if needed | |
| if mask.shape[0] != img_h or mask.shape[1] != img_w: | |
| mask = cv2.resize( | |
| mask.astype(np.float32), | |
| (img_w, img_h), | |
| interpolation=cv2.INTER_NEAREST, | |
| ) | |
| mask_bool = mask > 0.5 | |
| area = int(mask_bool.sum()) | |
| # Basic area filtering | |
| if area < min_area or area > max_area: | |
| continue | |
| bbox = self._mask_to_bbox(mask_bool) | |
| if bbox is None: | |
| continue | |
| # Shape quality filtering | |
| if not self._is_good_shape(mask_bool, bbox): | |
| continue | |
| conf = float(boxes.conf[i]) if boxes is not None and i < len(boxes.conf) else 0.0 | |
| segments.append( | |
| Segment( | |
| id=len(segments), | |
| mask=mask_bool, | |
| bbox=bbox, | |
| area=area, | |
| confidence=conf, | |
| ) | |
| ) | |
| # Remove overlapping segments | |
| segments = self._remove_overlaps(segments) | |
| # Sort by area descending | |
| segments.sort(key=lambda s: s.area, reverse=True) | |
| # Re-assign sequential ids | |
| for idx, seg in enumerate(segments): | |
| seg.id = idx | |
| return segments | |
| def _is_good_shape(self, mask: np.ndarray, bbox: list[int]) -> bool: | |
| """Filter out segments with poor shape characteristics. | |
| Args: | |
| mask: Boolean mask of the segment. | |
| bbox: Bounding box [x1, y1, x2, y2]. | |
| Returns: | |
| True if the segment has acceptable shape characteristics. | |
| """ | |
| x1, y1, x2, y2 = bbox | |
| bbox_width = x2 - x1 | |
| bbox_height = y2 - y1 | |
| # Skip very thin segments (likely artifacts or edges) | |
| if bbox_width < 5 or bbox_height < 5: | |
| return False | |
| # Calculate aspect ratio | |
| aspect_ratio = max(bbox_width, bbox_height) / max(min(bbox_width, bbox_height), 1) | |
| # Skip extremely elongated segments (likely edges or artifacts) | |
| if aspect_ratio > 10: | |
| return False | |
| # Calculate compactness (area / bbox_area) | |
| # Good segments should fill their bounding box reasonably | |
| bbox_area = bbox_width * bbox_height | |
| mask_area = np.sum(mask) | |
| compactness = mask_area / max(bbox_area, 1) | |
| # Skip very sparse segments | |
| if compactness < 0.15: | |
| return False | |
| return True | |
| def _remove_overlaps(self, segments: list[Segment], iou_threshold: float = 0.7) -> list[Segment]: | |
| """Remove highly overlapping segments, keeping the one with higher confidence. | |
| Args: | |
| segments: List of segments to filter. | |
| iou_threshold: IoU threshold above which segments are considered overlapping. | |
| Returns: | |
| Filtered list of segments with overlaps removed. | |
| """ | |
| if len(segments) <= 1: | |
| return segments | |
| # Sort by confidence (descending) to keep higher confidence segments | |
| segments_sorted = sorted(segments, key=lambda s: s.confidence, reverse=True) | |
| keep = [] | |
| for seg in segments_sorted: | |
| # Check if this segment overlaps significantly with any kept segment | |
| overlaps = False | |
| for kept_seg in keep: | |
| iou = self._calculate_iou(seg.mask, kept_seg.mask) | |
| if iou > iou_threshold: | |
| overlaps = True | |
| break | |
| if not overlaps: | |
| keep.append(seg) | |
| return keep | |
| def _calculate_iou(mask1: np.ndarray, mask2: np.ndarray) -> float: | |
| """Calculate Intersection over Union between two masks.""" | |
| intersection = np.logical_and(mask1, mask2).sum() | |
| union = np.logical_or(mask1, mask2).sum() | |
| if union == 0: | |
| return 0.0 | |
| return float(intersection) / float(union) | |
| def _mask_to_bbox(mask: np.ndarray) -> list[int] | None: | |
| """Compute bounding box [x1, y1, x2, y2] from a boolean mask.""" | |
| rows = np.any(mask, axis=1) | |
| cols = np.any(mask, axis=0) | |
| if not rows.any() or not cols.any(): | |
| return None | |
| y1, y2 = np.where(rows)[0][[0, -1]] | |
| x1, x2 = np.where(cols)[0][[0, -1]] | |
| return [int(x1), int(y1), int(x2 + 1), int(y2 + 1)] | |