wechat-screenshot-vision-algorithm 0.1.0__py3-none-any.whl

wechat_screenshot_vision_algorithm/__init__.py

@@ -0,0 +1,40 @@
+"""Public API for wechat_screenshot_vision_algorithm."""
+
+from wechat_screenshot_vision_algorithm._config import Platform, WeChatVersion, Profile
+
+from wechat_screenshot_vision_algorithm.algorithms.card_bbox import (
+    _compute_exact_card_bboxes,
+    BubbleBbox,
+    FavoriteLabelHit,
+    ThumbnailCard,
+    TrackedCard,
+    derive_cards,
+    drop_top_clamped_false_positive_cards,
+    bbox_to_metadata_list,
+    card_bounding_tuple,
+    card_overlaps_processed,
+    click_context_for_tap_thumbnail,
+    pick_first_unprocessed_card,
+    pick_next_unclicked_card,
+    y_interval_overlap_ratio,
+)
+
+__all__ = [
+    "Platform",
+    "WeChatVersion",
+    "Profile",
+    "ThumbnailCard",
+    "TrackedCard",
+    "FavoriteLabelHit",
+    "BubbleBbox",
+    "_compute_exact_card_bboxes",
+    "derive_cards",
+    "drop_top_clamped_false_positive_cards",
+    "bbox_to_metadata_list",
+    "card_bounding_tuple",
+    "card_overlaps_processed",
+    "click_context_for_tap_thumbnail",
+    "pick_first_unprocessed_card",
+    "pick_next_unclicked_card",
+    "y_interval_overlap_ratio",
+]

wechat_screenshot_vision_algorithm/_config.py

@@ -0,0 +1,61 @@
+"""Platform profile: parameter resolution by (platform, wechat_version)."""
+
+from dataclasses import dataclass
+from enum import Enum
+from importlib import import_module
+from pathlib import Path
+from typing import Any
+
+
+class Platform(str, Enum):
+    ANDROID = "android"
+    IOS = "ios"
+    HARMONY = "harmony"
+
+
+class WeChatVersion(str, Enum):
+    V8_0_69 = "8.0.69"
+
+
+#: All design constants are anchored to 1080x screen width (WeChat baseline).
+BASELINE_WIDTH = 1080
+
+_PROFILE_MODULES = {
+    Platform.ANDROID: "wechat_screenshot_vision_algorithm.profiles.android_wechat",
+    Platform.IOS: "wechat_screenshot_vision_algorithm.profiles.ios_wechat",
+    Platform.HARMONY: "wechat_screenshot_vision_algorithm.profiles.harmony_wechat",
+}
+
+_PACKAGE_ROOT = Path(__file__).resolve().parent
+
+
+@dataclass(frozen=True)
+class Profile:
+    """Parameter bundle for one (platform, wechat_version) combination.
+
+    Usage:
+        p = Profile(platform=Platform.ANDROID, wechat_version=WeChatVersion.V8_0_69)
+        templates = p.templates_dir  # Path to versioned PNG assets
+    """
+
+    platform: Platform
+    wechat_version: WeChatVersion
+
+    @property
+    def _profile_module(self) -> Any:
+        name = _PROFILE_MODULES[self.platform]
+        return import_module(name)
+
+    @property
+    def templates_dir(self) -> Path:
+        return _PACKAGE_ROOT / "templates" / "wechat" / self.platform.value / self.wechat_version.value
+
+    def __getattr__(self, name: str) -> Any:
+        """Delegate unknown attributes to the platform profile module."""
+        try:
+            return getattr(self._profile_module, name)
+        except AttributeError:
+            raise AttributeError(
+                f"Profile {self.platform.value}/{self.wechat_version.value} "
+                f"has no attribute '{name}'"
+            ) from None

wechat_screenshot_vision_algorithm/algorithms/avatar_column.py

@@ -0,0 +1,211 @@
+"""左列头像圆心检测（列表页 / 群聊页共用几何先验）。
+
+列表会话行与群聊发言人归因共用同一套 Hough + median_x 过滤 + y 向去重，
+产出 ``AvatarCentroid`` 供 ROI 守门与 PRD §5 纵向筒使用。
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+
+logger = logging.getLogger("wx_processor.left_avatar_column")
+
+try:
+    import cv2
+except ImportError:
+    cv2 = None  # type: ignore[assignment]
+
+from wechat_screenshot_vision_algorithm.algorithms.template_matching import (
+    CONVERSATION_LIST_AVATAR_MEDIAN_X_HALF_WIDTH_BASELINE,
+    CONVERSATION_LIST_AVATAR_ROI_LEFT_BASELINE,
+    CONVERSATION_LIST_AVATAR_ROI_RIGHT_BASELINE,
+)
+
+# 与列表页 Hough 对齐（1080 基准）；群聊左列头像物理尺寸与列表一致。
+_LIST_HOUGH_MIN_DIST_BASELINE = 117
+_LIST_HOUGH_MIN_R_BASELINE = 36
+_LIST_HOUGH_MAX_R_BASELINE = 80
+_LIST_YMIN_GAP_BASELINE = 118
+_LIST_SAME_ICON_DY_MAX_BASELINE = 18
+
+# 绑定：昵称行与头像锚点的最大竖直间距（px @ 1080）。
+NICKNAME_AVATAR_BIND_MAX_DY_BASELINE = 140
+
+
+@dataclass(frozen=True)
+class AvatarCentroid:
+    """左列头像圆（整图像素坐标）。"""
+
+    cx: int
+    cy: int
+    r: int
+
+    @property
+    def y_top(self) -> int:
+        return int(self.cy - self.r)
+
+
+@dataclass(frozen=True)
+class LeftAvatarColumnLayout:
+    centroids: tuple[AvatarCentroid, ...]
+    median_cx: float
+    column_x_half_width_px: int
+
+    @property
+    def empty(self) -> bool:
+        return not self.centroids
+
+
+def _sw(scale_w: float) -> float:
+    return max(float(scale_w), 1e-6)
+
+
+def detect_left_avatar_column_layout(
+    screen_bgr: np.ndarray,
+    scale_w: float,
+    *,
+    y_top: int = 0,
+    y_bottom_excl: Optional[int] = None,
+) -> LeftAvatarColumnLayout:
+    """在左列 ROI 内 Hough 圆 → median_x 过滤 → y 向去重（与列表页同参）。"""
+    sw = _sw(scale_w)
+    xt = int(round(float(CONVERSATION_LIST_AVATAR_MEDIAN_X_HALF_WIDTH_BASELINE) * sw))
+    empty = LeftAvatarColumnLayout(centroids=(), median_cx=0.0, column_x_half_width_px=xt)
+
+    if cv2 is None or screen_bgr is None or getattr(screen_bgr, "size", 0) == 0:
+        return empty
+
+    h_full, w_full = screen_bgr.shape[:2]
+    y1 = int(max(0, min(y_top, h_full - 1)))
+    y2 = int(y_bottom_excl if y_bottom_excl is not None else h_full)
+    y2 = int(max(y1 + 8, min(y2, h_full)))
+    x1_roi = int(max(0, round(CONVERSATION_LIST_AVATAR_ROI_LEFT_BASELINE * sw)))
+    x2_roi = int(min(w_full, round(CONVERSATION_LIST_AVATAR_ROI_RIGHT_BASELINE * sw)))
+    if y2 <= y1 + 40 or x2_roi <= x1_roi + 24:
+        return empty
+
+    crop = screen_bgr[y1:y2, x1_roi:x2_roi]
+    gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
+    gray_blur = cv2.GaussianBlur(gray, (5, 5), 0)
+    circles = cv2.HoughCircles(
+        gray_blur,
+        cv2.HOUGH_GRADIENT,
+        dp=1.25,
+        minDist=int(round(_LIST_HOUGH_MIN_DIST_BASELINE * sw)),
+        param1=100,
+        param2=22,
+        minRadius=max(8, int(round(_LIST_HOUGH_MIN_R_BASELINE * sw))),
+        maxRadius=int(round(_LIST_HOUGH_MAX_R_BASELINE * sw)),
+    )
+    if circles is None:
+        return empty
+
+    cand: list[tuple[float, float, float]] = []
+    for xi, yi, ri in circles[0]:
+        cand.append((float(ri), float(xi), float(yi)))
+
+    if not cand:
+        return empty
+
+    xs_crop = [c[1] for c in cand]
+    median_x_crop = float(np.median(np.array(xs_crop, dtype=np.float64)))
+    median_cx_full = float(x1_roi) + median_x_crop
+    filt = [c for c in cand if abs(c[1] - median_x_crop) <= xt]
+    if not filt:
+        return empty
+
+    filt.sort(key=lambda t: t[2])
+    ymin_gap = round(_LIST_YMIN_GAP_BASELINE * sw)
+    same_dy = float(_LIST_SAME_ICON_DY_MAX_BASELINE) * sw
+    stacked: list[tuple[float, float, float]] = []
+    for r, xc, yc in filt:
+        if not stacked:
+            stacked.append((r, xc, yc))
+            continue
+        pr, _px, py_prev = stacked[-1]
+        dy = float(yc) - float(py_prev)
+        if dy < ymin_gap:
+            if dy <= same_dy and r > pr:
+                stacked[-1] = (r, xc, yc)
+        else:
+            stacked.append((r, xc, yc))
+
+    centroids: list[AvatarCentroid] = []
+    for r, xc, yc in stacked:
+        cx_full = int(round(x1_roi + xc))
+        cy_full = int(round(y1 + yc))
+        ri = int(round(r))
+        centroids.append(AvatarCentroid(cx=cx_full, cy=cy_full, r=max(8, ri)))
+
+    return LeftAvatarColumnLayout(
+        centroids=tuple(centroids),
+        median_cx=median_cx_full,
+        column_x_half_width_px=xt,
+    )
+
+
+def find_avatar_anchor_for_nickname_bbox(
+    layout: LeftAvatarColumnLayout,
+    bbox_xyxy: tuple[int, int, int, int],
+    *,
+    scale_w: float,
+    max_dy_px: Optional[int] = None,
+) -> Optional[AvatarCentroid]:
+    """为昵称行找最近且在其上方的左列头像（几何绑定）。"""
+    if layout.empty:
+        return None
+    x1, y1, x2, y2 = bbox_xyxy
+    nick_cy = (y1 + y2) / 2.0
+    sw = _sw(scale_w)
+    max_dy = (
+        int(max_dy_px)
+        if max_dy_px is not None
+        else math.floor(NICKNAME_AVATAR_BIND_MAX_DY_BASELINE * sw)
+    )
+    nick_w = max(1, x2 - x1)
+    x_hi = float(layout.median_cx) + float(layout.column_x_half_width_px) + nick_w * 2
+
+    best: Optional[AvatarCentroid] = None
+    best_score = 1e18
+
+    for c in layout.centroids:
+        if float(c.cx) > x_hi:
+            continue
+        dy = float(y1) - float(c.cy)
+        if dy < -100:
+            continue
+        if dy >= max_dy:
+            continue
+        score = abs(float(nick_cy) - float(c.cy))
+        if score < best_score:
+            best_score = score
+            best = c
+
+    return best
+
+
+def nickname_bbox_in_avatar_column(
+    layout: LeftAvatarColumnLayout,
+    bbox_xyxy: tuple[int, int, int, int],
+    *,
+    original_width: int,
+    max_x1_ratio: float = 0.30,
+) -> bool:
+    """昵称行是否落在左列头像带（median_x 优先，无检出时回退 x1 比例）。"""
+    x1, _, x2, _ = bbox_xyxy
+    if layout.empty or layout.median_cx <= 0:
+        return float(x1) < float(original_width) * float(max_x1_ratio)
+    nick_w = max(1, x2 - x1)
+    lo = float(layout.median_cx) - float(layout.column_x_half_width_px)
+    hi = float(layout.median_cx) + float(layout.column_x_half_width_px) + nick_w * 2
+    return lo <= float(x1) <= hi
+
+
+def left_avatar_ytops(layout: LeftAvatarColumnLayout) -> list[int]:
+    """供 PRD §5：头像顶 y（圆心减半径）。"""
+    return [c.y_top for c in layout.centroids]

wechat_screenshot_vision_algorithm/algorithms/badge_detection.py

@@ -0,0 +1,275 @@
+"""Experimental list unread-badge heuristic: horizontal band × HSV red × bright glyphs.
+
+Used for calibration / offline scripts (not wired into scan_pinned_groups by default).
+Aligned with UX: badge sits on avatar left column — constrain x as fractions of screen width.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import cv2
+import numpy as np
+
+
+@dataclass(frozen=True)
+class WhiteOnRedBadgeCandidate:
+    """One candidate unread badge (~= disk-like red blob with white glyphs)."""
+
+    x: int
+    y: int
+    w: int
+    h: int
+    contour_area: int
+    bright_ratio: float
+    #: ``4*pi*area/perimeter²`` (~1 disk, ~0.78 square)
+    circularity: float
+    #: ``max(fitEllipse axes) / min(axes)`` (≈1 circle; >>1 elongated)
+    ellipse_axis_ratio: float
+    score: float
+
+
+def _dual_red_mask(hsv_roi: np.ndarray, *, sat_min: int, val_min: int, low_h_max: int) -> np.ndarray:
+    low_lo = np.array((0, sat_min, val_min), dtype=np.uint8)
+    low_hi = np.array((low_h_max, 255, 255), dtype=np.uint8)
+    high_lo = np.array((170, sat_min, val_min), dtype=np.uint8)
+    high_hi = np.array((180, 255, 255), dtype=np.uint8)
+    return cv2.bitwise_or(
+        cv2.inRange(hsv_roi, low_lo, low_hi),
+        cv2.inRange(hsv_roi, high_lo, high_hi),
+    )
+
+
+def _contour_circularity_pixels(contour: np.ndarray) -> float:
+    """Isoperimetric quotient; ~1 full disk, lower for dented/skinny blobs."""
+    area = abs(float(cv2.contourArea(contour)))
+    peri = float(cv2.arcLength(contour, True))
+    if peri < 1e-6:
+        return 0.0
+    return float((4.0 * math.pi * max(area, 1e-9)) / (peri * peri))
+
+
+def _ellipse_axis_ratio_or_none(contour: np.ndarray) -> float | None:
+    """``max(fitEllipse axis) / min(axis)``. ``None`` if ``len < 5``."""
+    if len(contour) < 5:
+        return None
+    (_, _), (maj, mn), _ = cv2.fitEllipse(contour)
+    mj, nn = float(maj), float(mn)
+    if mj < 1e-9 or nn < 1e-9:
+        return None
+    return float(max(mj, nn) / min(mj, nn))
+
+
+def count_badges_white_on_red(
+    screen_bgr: np.ndarray,
+    *,
+    width_frac_lo: float = 0.085,
+    width_frac_hi: float = 0.195,
+    vertical_y1: int = 270,
+    vertical_y2: int | None = None,
+    sat_min: int = 88,
+    val_min: int = 82,
+    low_h_max: int = 15,
+    min_area_px: int = 150,
+    max_area_px: int = 4000,
+    max_aspect: float = 1.95,
+    min_bright_ratio: float = 0.04,
+    white_gray_floor: int = 204,
+    morph_extra_dilate_iters: int = 1,
+    max_bbox_center_x_frac: float | None = 0.185,
+    min_bbox_center_y_abs: int | None = None,
+    max_bright_ratio: float | None = 0.17,
+) -> tuple[tuple[WhiteOnRedBadgeCandidate, ...], tuple[int, int, int, int]]:
+    """Find unread-style badges inside an x-percent band.
+
+    Steps (per user's scheme):
+      1) Restrict to horizontal band ``[width_frac_lo, width_frac_hi]`` of screen width
+         plus optional vertical clipping (conversation list bodies).
+      2) Dual-band HSV red mask + morphology.
+      3) Contour filter (area / aspect).
+      4) Reject blobs where the **filled contour** lacks enough luminous pixels
+         (**bright_ratio** = #(gray≥``white_gray_floor``∩red fill) / #red fill pixels).
+
+    Returns:
+        Sorted candidates top→bottom, and ``(roi_x1, roi_y1, roi_x2, roi_y2)``.
+    """
+    h_full, w_full = screen_bgr.shape[:2]
+    vx2 = vertical_y2 if vertical_y2 is not None else h_full - 120
+
+    roi_x1 = max(0, int(w_full * width_frac_lo))
+    roi_x2 = min(w_full, int(w_full * width_frac_hi))
+    roi_y1 = max(0, vertical_y1)
+    roi_y2 = min(h_full, vx2)
+
+    if roi_x2 <= roi_x1 or roi_y2 <= roi_y1:
+        return (), (0, 0, 0, 0)
+
+    roi_rgb = screen_bgr[roi_y1:roi_y2, roi_x1:roi_x2]
+    gray_roi = cv2.cvtColor(roi_rgb, cv2.COLOR_BGR2GRAY)
+
+    hsv = cv2.cvtColor(roi_rgb, cv2.COLOR_BGR2HSV)
+    mask = _dual_red_mask(hsv, sat_min=sat_min, val_min=val_min, low_h_max=low_h_max)
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
+    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
+    if morph_extra_dilate_iters > 0:
+        mask = cv2.dilate(mask, kernel, iterations=morph_extra_dilate_iters)
+
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    hits: list[WhiteOnRedBadgeCandidate] = []
+
+    for c in contours:
+        rx, ry, rw, rh = cv2.boundingRect(c)
+        area = float(cv2.contourArea(c))
+        if area < min_area_px or area > max_area_px:
+            continue
+        ls, ss = max(rw, rh), max(1, min(rw, rh))
+        if ls / ss > max_aspect:
+            continue
+
+        cx = rx + rw // 2
+        cy = ry + rh // 2
+        if max_bbox_center_x_frac is not None:
+            if (roi_x1 + cx) > int(w_full * max_bbox_center_x_frac):
+                continue
+        gc_y_abs = roi_y1 + cy
+        if min_bbox_center_y_abs is not None and gc_y_abs < min_bbox_center_y_abs:
+            continue
+
+        # White-ish digits: luminous pixels strictly inside filled red silhouette.
+        h_roi0, w_roi0 = gray_roi.shape
+        fill = np.zeros((h_roi0, w_roi0), dtype=np.uint8)
+        cv2.drawContours(fill, [c], 0, 255, -1)
+        bright_on_fill = np.count_nonzero(
+            np.logical_and(fill > 0, gray_roi >= np.uint8(white_gray_floor))
+        )
+        red_px = max(1, int(cv2.countNonZero(fill)))
+        bright_ratio = float(bright_on_fill) / float(red_px)
+
+        if bright_ratio < min_bright_ratio:
+            continue
+        if max_bright_ratio is not None and bright_ratio > max_bright_ratio:
+            continue
+
+        circ = _contour_circularity_pixels(c)
+        ear_raw = _ellipse_axis_ratio_or_none(c)
+        ear = float(ear_raw) if ear_raw is not None else 1.0
+
+        score = area * bright_ratio
+
+        gx = roi_x1 + rx
+        gy = roi_y1 + ry
+        hits.append(
+            WhiteOnRedBadgeCandidate(
+                x=int(gx),
+                y=int(gy),
+                w=int(rw),
+                h=int(rh),
+                contour_area=int(area),
+                bright_ratio=float(bright_ratio),
+                circularity=float(circ),
+                ellipse_axis_ratio=ear,
+                score=float(score),
+            )
+        )
+
+    hits.sort(key=lambda u: u.y)
+    merged = _dedupe_by_vertical_bucket(hits, bucket_px=95)
+    return tuple(merged), (roi_x1, roi_y1, roi_x2, roi_y2)
+
+
+def _dedupe_by_vertical_bucket(
+    hits: list[WhiteOnRedBadgeCandidate],
+    *,
+    bucket_px: float,
+) -> list[WhiteOnRedBadgeCandidate]:
+    """Merge multiple contours from the same list row (~one badge split)."""
+
+    if not hits:
+        return []
+
+    buckets: dict[int, list[WhiteOnRedBadgeCandidate]] = {}
+    for h in hits:
+        key = int((h.y + h.h // 2) // bucket_px)
+        buckets.setdefault(key, []).append(h)
+
+    merged: list[WhiteOnRedBadgeCandidate] = []
+    for grp in buckets.values():
+        grp.sort(key=lambda u: (-u.score, -u.contour_area))
+        merged.append(grp[0])
+    merged.sort(key=lambda u: u.y)
+    return merged
+
+
+def count_badges_per_row_roi(
+    screen_bgr: np.ndarray,
+    *,
+    scale_w: float | None = None,
+    min_bright_ratio: float = 0.04,
+    min_area_px: int = 260,
+    max_bright_ratio: float | None = 0.17,
+    **kw: object,
+) -> tuple[tuple[WhiteOnRedBadgeCandidate, ...], tuple[int, int, int, int]]:
+    """Same heuristic but matches :meth:`scan_pinned_groups` per-row avatar band.
+
+    Reuses DD row-split + ``badge_roi`` horizontal strip (stronger than one
+    flat width-% box).
+    """
+
+    del kw  # reserved for forwards-compat
+
+    from collector.driver import (
+        AVATAR_UNREAD_BADGE_LEFT_INSET_BASELINE,
+        AVATAR_UNREAD_BADGE_RIGHT_PAD_BASELINE,
+        CONVERSATION_AVATAR_RIGHT_EDGE_BASELINE,
+    )
+    from wechat_screenshot_vision_algorithm.algorithms.template_matching import split_conversation_rows
+
+    h_full, w_full = screen_bgr.shape[:2]
+    sw = scale_w if scale_w is not None else float(w_full) / 1080.0
+    row_rects = split_conversation_rows(h_full, sw)
+    avatar_edge = int(round(CONVERSATION_AVATAR_RIGHT_EDGE_BASELINE * sw))
+    ain = int(round(AVATAR_UNREAD_BADGE_LEFT_INSET_BASELINE * sw))
+    arp = int(round(AVATAR_UNREAD_BADGE_RIGHT_PAD_BASELINE * sw))
+
+    all_hits: list[WhiteOnRedBadgeCandidate] = []
+    for row_idx, (y1, y2) in enumerate(row_rects):
+        rh = y2 - y1
+        vy1 = y1 + max(1, int(rh * 0.04))
+        vy2 = y1 + int(rh * 0.58)
+        vy2 = max(vy1 + 8, vy2)
+        x1_bb = max(0, avatar_edge - ain)
+        x2_bb = min(w_full, avatar_edge + arp)
+        frac_lo = float(x1_bb) / float(w_full)
+        frac_hi = float(x2_bb) / float(w_full)
+        # Below the toolbar "微信(N)" reds (edb1a89f @1080-wide: ~282–294)
+        row_min_cy_abs: int | None = y1 + 18 if row_idx == 0 else None
+        cand, roi = count_badges_white_on_red(
+            screen_bgr,
+            width_frac_lo=frac_lo,
+            width_frac_hi=frac_hi,
+            vertical_y1=vy1,
+            vertical_y2=vy2,
+            min_bright_ratio=min_bright_ratio,
+            min_area_px=min_area_px,
+            max_bright_ratio=max_bright_ratio,
+            min_bbox_center_y_abs=row_min_cy_abs,
+        )
+        if cand:
+            all_hits.append(cand[0])
+
+    x_lo = max(0, int(round((CONVERSATION_AVATAR_RIGHT_EDGE_BASELINE - AVATAR_UNREAD_BADGE_LEFT_INSET_BASELINE) * sw)))
+    x_hi = min(
+        w_full,
+        int(round((CONVERSATION_AVATAR_RIGHT_EDGE_BASELINE + AVATAR_UNREAD_BADGE_RIGHT_PAD_BASELINE) * sw)),
+    )
+    y_lo = row_rects[0][0] if row_rects else 0
+    outer_roi = (x_lo, y_lo, x_hi, h_full)
+    return tuple(all_hits), outer_roi
+
+
+__all__ = [
+    "WhiteOnRedBadgeCandidate",
+    "count_badges_white_on_red",
+    "count_badges_per_row_roi",
+]