video-subtitle-remover/backend/tools/smart_detect.py

"""
智能水印检测：通过帧差分找静态区域，用SAM2精确分割。
"""
import sys
import os
import numpy as np
import cv2

SAM2_PATH = "/opt/sam2"
SAM2_CKPT = "/opt/sam2/checkpoints/sam2.1_hiera_tiny.pt"
SAM2_CFG  = "configs/sam2.1/sam2.1_hiera_t.yaml"


def detect_static_regions(video_path: str, sample_frames: int = 30, threshold: float = 20.0,
                          max_scan_seconds: float = None):
    """
    饱和度签名 + 跨帧对比法检测水印/Logo：
    1. 采样多帧，通过高饱和度像素比例判断哪些帧含有logo（logo通常色彩鲜艳）
    2. 在"含logo帧"之间比较，logo像素在多帧中出现在同一位置 → 差值极低
    3. 能检测周期性出现/动态水印，无需logo持续存在于全部帧
    """
    cap = cv2.VideoCapture(video_path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    duration_s = total / fps
    num_samples = max(10, min(24, int(duration_s // 8)))
    positions = [min(int((i + 0.5) * total / num_samples), total - 1)
                 for i in range(num_samples)]

    bgr_frames = []
    gray_frames = []
    for pos in positions:
        cap.set(cv2.CAP_PROP_POS_FRAMES, pos)
        ret, frame = cap.read()
        if ret:
            bgr_frames.append(frame)
            gray_frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY).astype(np.float32))
    cap.release()

    print(f"[SAM2-detect] 采样 {len(bgr_frames)} 帧 (均匀分布于{duration_s:.1f}s视频)", flush=True)

    if len(bgr_frames) < 4:
        print(f"[SAM2-detect] 帧数不足4帧，返回空", flush=True)
        return []

    # ── 阶段1：通过高饱和度识别含logo帧 ──────────────────────────────────
    # 水印/logo通常色彩比视频自然内容更鲜艳（高饱和度）
    # 在帧的上半部分（logo常出现区域）检测高饱和度像素数
    top_h = h // 5  # 检查帧顶部 20%
    logo_sat_thresh = 200   # 饱和度 > 200 认为是logo颜色
    logo_min_pixels = top_h * w * 0.003  # 至少 0.3% 的上方区域为logo颜色

    logo_frames_gray = []
    for i, (bgr, gray) in enumerate(zip(bgr_frames, gray_frames)):
        hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
        high_sat = int(np.sum(hsv[:top_h, :, 1] > logo_sat_thresh))
        pos_s = positions[i] / fps
        if high_sat >= logo_min_pixels:
            logo_frames_gray.append(gray)
            print(f"[SAM2-detect]   ✓ {pos_s:.1f}s: 高饱和={high_sat}px (含logo)", flush=True)
        else:
            print(f"[SAM2-detect]   ✗ {pos_s:.1f}s: 高饱和={high_sat}px (无logo)", flush=True)

    print(f"[SAM2-detect] logo含帧: {len(logo_frames_gray)}/{len(bgr_frames)}", flush=True)

    logo_bgr_frames = [bgr_frames[i] for i, g in enumerate(gray_frames)
                       if any(np.array_equal(g, l) for l in logo_frames_gray)]
    nlogo_bgr_frames = [bgr_frames[i] for i, g in enumerate(gray_frames)
                        if not any(np.array_equal(g, l) for l in logo_frames_gray)]

    if len(logo_frames_gray) < 2:
        print(f"[SAM2-detect] logo帧不足，回退到时序std法", flush=True)
        stack = np.stack(gray_frames, axis=0)
        std_map = np.std(stack, axis=0)
        static_mask = (std_map < 8.0).astype(np.uint8) * 255
    else:
        # 在logo帧中统计每像素平均饱和度
        sat_acc = np.zeros((h, w), dtype=np.float32)
        for bgr in logo_bgr_frames:
            hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
            sat_acc += (hsv[:, :, 1].astype(np.float32) / 255.0)
        sat_mean_logo = sat_acc / len(logo_bgr_frames)

        if nlogo_bgr_frames:
            # ── 有无logo帧对比：准确检测logo区域 ───────────────────────────
            # Logo像素: logo帧高饱和 + 无logo帧低饱和（差值大）
            sat_acc_nl = np.zeros((h, w), dtype=np.float32)
            for bgr in nlogo_bgr_frames[:6]:
                hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV)
                sat_acc_nl += (hsv[:, :, 1].astype(np.float32) / 255.0)
            sat_mean_nlogo = sat_acc_nl / min(len(nlogo_bgr_frames), 6)
            sat_diff = sat_mean_logo - sat_mean_nlogo

            sat_thresh = min(0.6, np.percentile(sat_mean_logo, 92))
            diff_thresh = max(0.25, sat_diff.mean() + sat_diff.std())
            logo_pixel_mask = ((sat_mean_logo >= sat_thresh) & (sat_diff >= diff_thresh)).astype(np.uint8) * 255

            print(f"[SAM2-detect] logo帧平均饱和: max={sat_mean_logo.max():.3f}", flush=True)
            print(f"[SAM2-detect] 饱和差值: max={sat_diff.max():.3f} mean={sat_diff.mean():.3f}", flush=True)
            print(f"[SAM2-detect] sat_thresh={sat_thresh:.3f} diff_thresh={diff_thresh:.3f}, 候选={np.sum(logo_pixel_mask > 0)}", flush=True)
        else:
            # ── 无非logo帧（短视频/预览：logo全程可见）──────────────────────
            # 只在帧的边缘区域（顶部30%）寻找高饱和logo
            # 避免把视频中央的自然内容误报为logo
            border_mask = np.zeros((h, w), dtype=np.uint8)
            border_mask[:h * 3 // 10, :] = 1  # 顶部 30%

            # 取最高饱和的4%像素（极鲜艳 = logo文字/图形）
            sat_thresh = max(0.6, np.percentile(sat_mean_logo, 96))
            logo_pixel_mask = ((sat_mean_logo >= sat_thresh).astype(np.uint8) * border_mask * 255)

            print(f"[SAM2-detect] 无非logo帧，限定边缘区域搜索", flush=True)
            print(f"[SAM2-detect] sat_thresh={sat_thresh:.3f}, 候选={np.sum(logo_pixel_mask > 0)}", flush=True)

        static_mask = logo_pixel_mask

    # 仅排除纯黑区域
    mean_gray = np.mean(np.stack(gray_frames, axis=0), axis=0)
    content_mask = (mean_gray > 8).astype(np.uint8) * 255
    static_mask = cv2.bitwise_and(static_mask, content_mask)
    print(f"[SAM2-detect] 过滤纯黑后候选像素数={np.sum(static_mask > 0)}", flush=True)

    # 形态学处理：去噪 + 膨胀
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    static_mask = cv2.morphologyEx(static_mask, cv2.MORPH_OPEN, kernel)
    static_mask = cv2.dilate(static_mask, kernel, iterations=2)

    # 找轮廓
    contours, _ = cv2.findContours(static_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    regions = []
    min_area = (w * h) * 0.0005  # 至少占画面 0.05%
    max_area = (w * h) * 0.25    # 最多占画面 25%

    print(f"[SAM2-detect] 找到 {len(contours)} 个轮廓，min_area={min_area:.0f} max_area={max_area:.0f}", flush=True)
    for cnt in contours:
        area = cv2.contourArea(cnt)
        x, y, bw, bh = cv2.boundingRect(cnt)
        ratio = bw / bh if bh > 0 else 0
        if min_area < area < max_area:
            if 0.1 < ratio < 20:
                regions.append((x, y, x + bw, y + bh))
                print(f"[SAM2-detect]   ✓ 区域 ({x},{y})-({x+bw},{y+bh}) 面积={area:.0f} 比例={ratio:.2f}", flush=True)
            else:
                print(f"[SAM2-detect]   ✗ 比例过滤: ({x},{y})-({x+bw},{y+bh}) 面积={area:.0f} 比例={ratio:.2f}", flush=True)
        else:
            print(f"[SAM2-detect]   ✗ 面积过滤: ({x},{y})-({x+bw},{y+bh}) 面积={area:.0f}", flush=True)

    # 若检测到多个位于帧顶部的区域，合并成一个大框覆盖整个logo
    # logo通常有高饱和的亮色区域（检测到）和低饱和的暗色区域（未检测到）
    # 合并+适当扩展可确保完整覆盖
    if len(regions) >= 2:
        top_regions = [r for r in regions if r[1] < h * 0.3]  # 只合并帧顶30%内的区域
        if len(top_regions) >= 2:
            x1 = max(0, min(r[0] for r in top_regions) - 30)
            y1 = max(0, min(r[1] for r in top_regions) - 15)
            x2 = min(w, max(r[2] for r in top_regions) + 30)
            y2 = min(h, max(r[3] for r in top_regions) + 20)
            other = [r for r in regions if r[1] >= h * 0.3]
            merged = [(x1, y1, x2, y2)] + other
            print(f"[SAM2-detect] 合并顶部 {len(top_regions)} 个区域为 ({x1},{y1})-({x2},{y2})", flush=True)
            regions = merged

    print(f"[SAM2-detect] detect_static_regions 返回 {len(regions)} 个区域", flush=True)
    return regions


def refine_with_sam2(image: np.ndarray, boxes: list, gpu_id: int = 0):
    """
    用 SAM2 对检测到的区域精确分割，返回合并后的 mask。
    image: BGR numpy array
    boxes: [(x1,y1,x2,y2), ...]
    返回 binary mask (H, W) uint8
    """
    if not boxes:
        return None

    try:
        if SAM2_PATH not in sys.path:
            sys.path.insert(0, SAM2_PATH)

        import torch
        from sam2.build_sam import build_sam2
        from sam2.sam2_image_predictor import SAM2ImagePredictor

        device = f"cuda:{gpu_id}" if torch.cuda.is_available() else "cpu"
        model = build_sam2(SAM2_CFG, SAM2_CKPT, device=device)
        predictor = SAM2ImagePredictor(model)

        rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        predictor.set_image(rgb)

        h, w = image.shape[:2]
        combined_mask = np.zeros((h, w), dtype=np.uint8)

        for (x1, y1, x2, y2) in boxes:
            box = np.array([x1, y1, x2, y2])
            masks, scores, _ = predictor.predict(
                box=box[None],
                multimask_output=False,
            )
            if masks is not None and len(masks) > 0:
                m = masks[0].astype(np.uint8) * 255
                combined_mask = cv2.bitwise_or(combined_mask, m)

        return combined_mask

    except Exception as e:
        print(f"[SAM2] 分割失败: {e}", flush=True)
        return None


def smart_detect_watermark(video_path: str, sample_frames: int = 30, gpu_id: int = 0):
    """
    完整流程：帧差分（前30秒每秒一帧）→ SAM2 精确分割。
    限制扫描前30秒：H.264压缩导致长视频里即使静态logo的像素std也会随时间累积超过阈值。
    返回 [(x1, y1, x2, y2), ...] 像素坐标列表。
    """
    print(f"[SAM2] smart_detect_watermark 开始: {video_path}", flush=True)
    # 1. 静态区域检测：跨时间段帧对比法（全视频）
    boxes = detect_static_regions(video_path, sample_frames)
    print(f"[SAM2] detect_static_regions 返回 {len(boxes)} 个候选区域: {boxes}", flush=True)
    if not boxes:
        return []

    # 2. 读取代表帧用于 SAM2
    cap = cv2.VideoCapture(video_path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.set(cv2.CAP_PROP_POS_FRAMES, total // 2)
    ret, frame = cap.read()
    cap.release()
    if not ret:
        return boxes  # SAM2 失败时直接返回粗检测结果

    # 3. SAM2 精确分割
    mask = refine_with_sam2(frame, boxes, gpu_id=gpu_id)
    if mask is None:
        return boxes

    # 4. 从 mask 提取最终 bounding boxes
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    refined = []
    for cnt in contours:
        if cv2.contourArea(cnt) > 100:
            x, y, bw, bh = cv2.boundingRect(cnt)
            refined.append((x, y, x + bw, y + bh))

    result = refined if refined else boxes
    print(f"[SAM2] smart_detect_watermark 最终返回 {len(result)} 个区域: {result}", flush=True)
    return result