Create sam_segment.py

sam_segment.py

@@ -0,0 +1,136 @@
+import numpy as np
+import cv2
+from ultralytics import FastSAM
+import torch
+import gc
+
+# 定义可用的模型
+MODELS = {
+    "small": "/disk2/models/FastSAM-s.pt",
+    "large": "/disk2/models/FastSAM-x.pt"
+}
+
+def clear_gpu_memory():
+    """
+    清理GPU显存
+    """
+    gc.collect()  # 清理Python的垃圾收集器
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()  # 清空PyTorch的CUDA缓存
+        torch.cuda.ipc_collect()  # 收集CUDA IPC内存
+
+def get_model(model_size: str = "large"):
+    """
+    获取指定大小的模型
+    """
+    if model_size not in MODELS:
+        raise ValueError(f"Invalid model size. Available sizes: {list(MODELS.keys())}")
+    
+    try:
+        return FastSAM(MODELS[model_size])
+    except Exception as e:
+        raise RuntimeError(f"Failed to load model: {str(e)}")
+
+def mask_to_points(mask: np.ndarray) -> list:
+    """
+    Convert mask to a list of contour points
+    
+    Args:
+        mask: 2D numpy array representing the mask
+    
+    Returns:
+        list: Flattened list of points [x1, y1, x2, y2, ...]
+    """
+    # Convert mask to uint8 type
+    mask_uint8 = mask.astype(np.uint8) * 255
+    # Find contours
+    contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    
+    if not contours:
+        return []
+    
+    # Get the largest contour
+    contour = max(contours, key=cv2.contourArea)
+    # Convert contour points to flattened list
+    points = []
+    for point in contour:
+        points.extend([float(point[0][0]), float(point[0][1])])
+    return points
+
+def segment_image_with_prompt(
+    image: np.ndarray,
+    model_size: str = "large",
+    conf: float = 0.4,
+    iou: float = 0.9,
+    bboxes: list = None,
+    points: list = None,
+    labels: list = None,
+    texts: str = None
+):
+    """
+    带提示的图像分割函数
+    
+    Args:
+        image: 输入图像 (numpy.ndarray)
+        model_size: 模型大小 ("small" 或 "large")
+        conf: 置信度阈值
+        iou: IoU 阈值
+        bboxes: 边界框列表 [x1, y1, x2, y2, ...]
+        points: 点列表 [[x1, y1], [x2, y2], ...]
+        labels: 点对应的标签列表 [0, 1, ...]
+        texts: 文本提示
+    """
+    try:
+        if image is None:
+            raise ValueError("Invalid image input")
+        
+        # 获取模型并执行分割
+        model = get_model(model_size)
+        
+        # 准备模型参数
+        model_args = {
+            "retina_masks": True,
+            "conf": conf,
+            "iou": iou
+        }
+        
+        # 添加提示参数
+        if bboxes is not None:
+            model_args["bboxes"] = bboxes
+        if points is not None and labels is not None:
+            model_args["points"] = points
+            model_args["labels"] = labels
+        if texts is not None:
+            model_args["texts"] = texts
+            
+        # 执行分割
+        everything_results = model(image, **model_args)
+        
+        # 处理分割结果
+        segments = []
+        if everything_results and len(everything_results) > 0:
+            result = everything_results[0]
+            if hasattr(result, 'masks') and result.masks is not None:
+                masks = result.masks.data.cpu().numpy()
+                
+                for mask in masks:
+                    points = mask_to_points(mask)
+                    if points:
+                        segments.append(points)
+        
+        # 清理模型和GPU内存
+        del model
+        del everything_results
+        if hasattr(result, 'masks'):
+            del result.masks
+        del result
+        clear_gpu_memory()
+        
+        return {
+            "total_segments": len(segments),
+            "segments": segments
+        }
+    except Exception as e:
+        # 确保发生错误时也清理内存
+        clear_gpu_memory()
+        raise RuntimeError(f"Error processing image: {str(e)}")