Blame for infer.py

#!/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()