#!/usr/bin/env python3 """ SAM2 ONNX 独立推理脚本 (sam2-hiera-base-plus) 依赖: numpy, opencv-python, onnxruntime 无需 torch / CUDA 用法: # 单点分割 python infer.py --image photo.jpg --point 500 400 # 多点分割 python infer.py --image photo.jpg --point 500 400 --point 600 300 # bbox 中心点分割 python infer.py --image photo.jpg --bbox 100 200 500 600 # 保存结果 python infer.py --image photo.jpg --point 500 400 --save mask.png """ import argparse import json import os import sys import cv2 import numpy as np SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) VE_ONNX = os.path.join(SCRIPT_DIR, "vision_encoder.onnx") MD_ONNX = os.path.join(SCRIPT_DIR, "mask_decoder.onnx") CONV_S0_NPZ = os.path.join(SCRIPT_DIR, "conv_s0.npz") CONV_S1_NPZ = os.path.join(SCRIPT_DIR, "conv_s1.npz") PE_NPZ = os.path.join(SCRIPT_DIR, "prompt_encoder.npz") PROC_CONFIG = os.path.join(SCRIPT_DIR, "preprocessor_config.json") INPUT_SIZE = 1024 PE_INPUT_SIZE = 1024 # ============================================================================ # 图像预处理 # ============================================================================ def preprocess_image(image_rgb: np.ndarray): """ SAM2 图像预处理: resize longest edge to 1024, pad to 1024×1024, normalize 返回: pixel_values (1, 3, 1024, 1024), (orig_h, orig_w), scale """ h, w = image_rgb.shape[:2] scale = INPUT_SIZE / max(h, w) new_h, new_w = int(h * scale), int(w * scale) resized = cv2.resize(image_rgb, (new_w, new_h), interpolation=cv2.INTER_AREA) # Normalize (ImageNet-like, from SAM2 processor) img = resized.astype(np.float32) / 255.0 mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) std = np.array([0.229, 0.224, 0.225], dtype=np.float32) img = (img - mean) / std # Pad to INPUT_SIZE × INPUT_SIZE padded = np.zeros((INPUT_SIZE, INPUT_SIZE, 3), dtype=np.float32) padded[:new_h, :new_w] = img # HWC → CHW → NCHW pixel_values = padded.transpose(2, 0, 1)[np.newaxis] return pixel_values, (h, w), scale # ============================================================================ # Conv2d 推理 (conv_s0, conv_s1) # ============================================================================ def conv2d_forward(x, weight, bias): """纯 numpy Conv2d(1×1) 实现""" # weight: (out_ch, in_ch, 1, 1), x: (1, in_ch, H, W) B, C, H, W = x.shape out_ch = weight.shape[0] # 1×1 conv = matrix multiply on channel dimension x_flat = x.reshape(B, C, -1) # (B, C, H*W) w_flat = weight.reshape(out_ch, C) # (out_ch, C) out = np.einsum('ij,bjk->bik', w_flat, x_flat) + bias.reshape(1, -1, 1) return out.reshape(B, out_ch, H, W) # ============================================================================ # Prompt Encoder (纯 numpy 实现) # ============================================================================ class PromptEncoder: """SAM2 Prompt Encoder — 纯 numpy 实现""" def __init__(self, pe_weights_path): state = np.load(pe_weights_path) # Sinusoidal positional embedding matrix: (2, 128) self.pos_embed = state["shared_embedding.positional_embedding"] # Point type embeddings self.point_embed_weight = state["point_embed.weight"] # (4, 256) self.not_a_point_weight = state["not_a_point_embed.weight"] # (1, 256) self.no_mask_weight = state["no_mask_embed.weight"] # (1, 256) def _embed_points(self, coords): """ Sinusoidal positional encoding for 2D coordinates. coords: (N, n_points, 2) — in PE coordinate space (0 to PE_INPUT_SIZE) """ # Shift to center of pixel coords = coords + 0.5 # Normalize to [0, 1] coords[..., 0] = coords[..., 0] / PE_INPUT_SIZE coords[..., 1] = coords[..., 1] / PE_INPUT_SIZE # Scale to [-1, 1] coords = 2 * coords - 1 # Apply positional embedding matrix # coords: (N, P, 2) @ pos_embed: (2, D) → (N, P, D) embedded = coords @ self.pos_embed # Scale by 2π embedded = 2 * np.pi * embedded # Concat sin and cos sin_emb = np.sin(embedded) # (N, P, D) cos_emb = np.cos(embedded) # (N, P, D) return np.concatenate([sin_emb, cos_emb], axis=-1) # (N, P, 2D) def encode_points(self, points, labels): """ Encode point prompts. points: (B, n_prompts, n_points, 2) — raw pixel coordinates labels: (B, n_prompts, n_points) — 1=foreground, 0=background, -1=not_a_point Returns: sparse_emb (B, n_prompts, n_points, embed_dim), dense_emb (B, embed_dim, 64, 64) """ B, n_prompts, n_points, _ = points.shape embed_dim = self.pos_embed.shape[1] * 2 # 2D = sin + cos # Flatten for encoding coords_flat = points.reshape(B, -1, 2) # (B, n_prompts*n_points, 2) labels_flat = labels.reshape(B, -1) # (B, n_prompts*n_points) # Sinusoidal position encoding pos_enc = self._embed_points(coords_flat) # (B, n_total, embed_dim) # Add type embeddings based on labels # label=1 → point_embed (foreground) # label=0 → not_a_point (should not appear normally) # label=-1 → not_a_point_embed fg_mask = (labels_flat == 1).astype(np.float32)[..., np.newaxis] bg_mask = (labels_flat == 0).astype(np.float32)[..., np.newaxis] na_mask = (labels_flat == -1).astype(np.float32)[..., np.newaxis] # point_embed.weight: (4, 256) — nn.Embedding(4, 256) # label=0 → row 0 (background), label=1 → row 1 (foreground) fg_emb = self.point_embed_weight[1] # (256,) — label=1 foreground bg_emb = self.point_embed_weight[0] # (256,) — label=0 background na_emb = self.not_a_point_weight[0] # (256,) — label=-1 not_a_point type_emb = (fg_mask * fg_emb + bg_mask * bg_emb + na_mask * na_emb) sparse_emb = pos_enc + type_emb # (B, n_total, embed_dim) # Reshape: (B, n_prompts, n_points, embed_dim) → (B, n_prompts, n_points*embed_dim) sparse_emb = sparse_emb.reshape(B, n_prompts, n_points, embed_dim) # For the ONNX mask decoder, sparse_prompt_emb shape is (B, n_prompts, n_points, embed_dim) # and dense_prompt_emb shape is (B, embed_dim, 64, 64) # Dense embedding: start with no_mask_embed, broadcast to spatial size dense_emb = np.broadcast_to( self.no_mask_weight.reshape(1, embed_dim, 1, 1), (B, embed_dim, 64, 64) ).copy() return sparse_emb.astype(np.float32), dense_emb.astype(np.float32) # ============================================================================ # 主推理流程 # ============================================================================ class SAM2Segmenter: """SAM2 独立分割器 — 使用 ONNX VE + MD""" def __init__(self, providers=None): import onnxruntime as ort if providers is None: providers = ["CPUExecutionProvider"] self.ve_sess = ort.InferenceSession(VE_ONNX, providers=providers) self.md_sess = ort.InferenceSession(MD_ONNX, providers=providers) # Load conv weights from .npz (no torch dependency) s0 = np.load(CONV_S0_NPZ) s1 = np.load(CONV_S1_NPZ) self.conv_s0_w = s0["weight"] self.conv_s0_b = s0["bias"] self.conv_s1_w = s1["weight"] self.conv_s1_b = s1["bias"] # Load prompt encoder self.pe = PromptEncoder(PE_NPZ) def encode_image(self, image_rgb): """运行 Vision Encoder, 返回图像特征缓存""" pixel_values, (orig_h, orig_w), scale = preprocess_image(image_rgb) ve_out = self.ve_sess.run(None, {"pixel_values": pixel_values}) # Outputs: last_hidden_state, fpn0, fpn1, fpn2, pos0, pos1, pos2 _, fpn0, fpn1, fpn2, _, _, pos2 = ve_out # Apply conv_s0 / conv_s1 hres0 = conv2d_forward(fpn0, self.conv_s0_w, self.conv_s0_b) hres1 = conv2d_forward(fpn1, self.conv_s1_w, self.conv_s1_b) return { "image_embeddings": fpn2, "image_pos_enc": pos2, "high_res_feat_0": hres0, "high_res_feat_1": hres1, "orig_h": orig_h, "orig_w": orig_w, "scale": scale, } def _scale_points(self, ve_cache, points_xy): """将原始图像像素坐标缩放到 1024×1024 预处理空间""" scale = ve_cache["scale"] return [(float(x) * scale, float(y) * scale) for x, y in points_xy] def segment_point(self, ve_cache, point_xy, label=1): """ 使用点 prompt 分割. point_xy: (x, y) 原始图像像素坐标 label: 1=前景, 0=背景 注意: ONNX mask_decoder 固定 n_points=2, 单点自动填充 not-a-point """ scaled = self._scale_points(ve_cache, [point_xy]) x, y = scaled[0] # 填充第二个点为 not-a-point (label=-1) points = np.array([[[[x, y], [0.0, 0.0]]]], dtype=np.float32) # (1, 1, 2, 2) labels = np.array([[[label, -1]]], dtype=np.int64) # (1, 1, 2) return self._run_decoder(ve_cache, points, labels) def segment_points(self, ve_cache, points_xy, labels_list): """ 使用多个点 prompt 分割. 每次取 2 个点运行, 合并 mask. points_xy: list of (x, y) — 原始图像像素坐标 labels_list: list of int (1=fg, 0=bg) """ scaled = self._scale_points(ve_cache, points_xy) if len(scaled) == 0: raise ValueError("No points provided") # 分组: 每 2 个点一组, 不足 2 个则填充 not-a-point combined_mask = None best_iou = 0.0 total_pixels = 0 for i in range(0, len(scaled), 2): chunk_pts = scaled[i:i+2] chunk_lbl = labels_list[i:i+2] if len(chunk_pts) == 1: chunk_pts.append((0.0, 0.0)) chunk_lbl.append(-1) points = np.array([[list(p) for p in chunk_pts]], dtype=np.float32) # (1, 1, 2, 2) labels = np.array([[chunk_lbl]], dtype=np.int64) # (1, 1, 2) result = self._run_decoder(ve_cache, points, labels) if combined_mask is None: combined_mask = result["mask"] else: combined_mask = np.maximum(combined_mask, result["mask"]) best_iou = max(best_iou, result["iou"]) total_pixels += result["pixels"] return { "mask": combined_mask, "iou": best_iou, "pixels": total_pixels, "raw_mask": None, } def segment_bbox_center(self, ve_cache, bbox_xyxy): """使用 bbox 中心点作为 prompt 分割 (bbox 为原始图像坐标)""" x1, y1, x2, y2 = bbox_xyxy cx = (x1 + x2) / 2.0 cy = (y1 + y2) / 2.0 return self.segment_point(ve_cache, (cx, cy), label=1) def _run_decoder(self, ve_cache, points, labels): """运行 Prompt Encoder + Mask Decoder""" sparse_emb, dense_emb = self.pe.encode_points(points, labels) md_out = self.md_sess.run(None, { "image_embeddings": ve_cache["image_embeddings"], "image_pos_enc": ve_cache["image_pos_enc"], "sparse_prompt_emb": sparse_emb, "dense_prompt_emb": dense_emb, "high_res_feat_0": ve_cache["high_res_feat_0"], "high_res_feat_1": ve_cache["high_res_feat_1"], }) iou_scores, pred_masks, obj_score = md_out # Pick best mask from 3 candidates (multimask_output=True) best_idx = np.argmax(iou_scores[0, 0]) best_iou = float(iou_scores[0, 0, best_idx]) mask = pred_masks[0, 0, best_idx] # (256, 256) # Upsample to original image size orig_h, orig_w = ve_cache["orig_h"], ve_cache["orig_w"] mask_resized = cv2.resize(mask, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR) binary_mask = (mask_resized > 0.0).astype(np.uint8) return { "mask": binary_mask, "iou": best_iou, "pixels": int(binary_mask.sum()), "raw_mask": mask, } # ============================================================================ # 可视化 # ============================================================================ def overlay_mask(image_rgb, mask, color=(0, 255, 0), alpha=0.5): """将 mask 叠加到图像上""" overlay = image_rgb.copy().astype(np.float32) color_arr = np.array(color, dtype=np.float32) overlay[mask > 0] = overlay[mask > 0] * (1 - alpha) + color_arr * alpha return np.clip(overlay, 0, 255).astype(np.uint8) # ============================================================================ # Main # ============================================================================ def main(): parser = argparse.ArgumentParser(description="SAM2 ONNX 独立分割推理") parser.add_argument("--image", required=True, help="输入图像路径") parser.add_argument("--point", nargs=2, type=float, action="append", metavar=("X", "Y"), help="点 prompt (可多次指定)") parser.add_argument("--bbox", nargs=4, type=float, metavar=("X1", "Y1", "X2", "Y2"), help="bbox (使用中心点)") parser.add_argument("--save", default=None, help="保存结果图像路径") parser.add_argument("--save-mask", default=None, help="保存纯 mask 图像路径") parser.add_argument("--providers", nargs="+", default=["CPUExecutionProvider"]) args = parser.parse_args() # 加载图像 image_bgr = cv2.imread(args.image) if image_bgr is None: raise FileNotFoundError(f"Image not found: {args.image}") image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB) h, w = image_rgb.shape[:2] # 初始化分割器 seg = SAM2Segmenter(args.providers) # 编码图像 (VE — 只需一次) print(f"Encoding image ({w}x{h})...") ve_cache = seg.encode_image(image_rgb) # 分割 if args.point: for i, (x, y) in enumerate(args.point): result = seg.segment_point(ve_cache, (x, y), label=1) print(f" Point ({x}, {y}): IoU={result['iou']:.4f}, pixels={result['pixels']}") if i == 0: combined_mask = result["mask"] else: combined_mask = np.maximum(combined_mask, result["mask"]) elif args.bbox: result = seg.segment_bbox_center(ve_cache, args.bbox) print(f" BBox center: IoU={result['iou']:.4f}, pixels={result['pixels']}") combined_mask = result["mask"] else: print("Error: specify --point or --bbox") sys.exit(1) # 保存结果 if args.save: result_img = overlay_mask(image_rgb, combined_mask) cv2.imwrite(args.save, cv2.cvtColor(result_img, cv2.COLOR_RGB2BGR)) print(f"Saved result to: {args.save}") if args.save_mask: cv2.imwrite(args.save_mask, combined_mask * 255) print(f"Saved mask to: {args.save_mask}") if __name__ == "__main__": main()