screenshot-vision-algorithm 0.3.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- screenshot_vision_algorithm/__init__.py +48 -0
- screenshot_vision_algorithm/_config.py +61 -0
- screenshot_vision_algorithm/android/__init__.py +1 -0
- screenshot_vision_algorithm/android/wechat/__init__.py +1 -0
- screenshot_vision_algorithm/android/wechat/algorithms/__init__.py +0 -0
- screenshot_vision_algorithm/android/wechat/algorithms/avatar_column.py +209 -0
- screenshot_vision_algorithm/android/wechat/algorithms/badge_detection.py +275 -0
- screenshot_vision_algorithm/android/wechat/algorithms/card_bbox.py +1000 -0
- screenshot_vision_algorithm/android/wechat/algorithms/phash_utils.py +267 -0
- screenshot_vision_algorithm/android/wechat/algorithms/speaker_band.py +290 -0
- screenshot_vision_algorithm/android/wechat/algorithms/template_matching.py +2163 -0
- screenshot_vision_algorithm/android/wechat/algorithms/title_ocr.py +143 -0
- screenshot_vision_algorithm/android/wechat/merge/__init__.py +0 -0
- screenshot_vision_algorithm/android/wechat/merge/multipage.py +157 -0
- screenshot_vision_algorithm/android/wechat/ocr/__init__.py +0 -0
- screenshot_vision_algorithm/android/wechat/ocr/avatar_guard.py +434 -0
- screenshot_vision_algorithm/android/wechat/ocr/badge_ocr.py +232 -0
- screenshot_vision_algorithm/android/wechat/ocr/nickname_binding.py +1888 -0
- screenshot_vision_algorithm/android/wechat/ocr/text_ocr_adapter.py +625 -0
- screenshot_vision_algorithm/android/wechat/profiles/__init__.py +0 -0
- screenshot_vision_algorithm/android/wechat/profiles/android.py +53 -0
- screenshot_vision_algorithm/android/wechat/profiles/harmony.py +10 -0
- screenshot_vision_algorithm/android/wechat/profiles/ios.py +53 -0
- screenshot_vision_algorithm/android/wechat/templates/android/8.0.69/chat_back_chevron.png +0 -0
- screenshot_vision_algorithm/android/wechat/templates/android/8.0.69/chat_input_emoji_smile.png +0 -0
- screenshot_vision_algorithm/android/wechat/templates/android/8.0.69/chat_input_plus.png +0 -0
- screenshot_vision_algorithm/android/wechat/templates/android/8.0.69/chat_input_voice.png +0 -0
- screenshot_vision_algorithm/android/wechat/templates/android/8.0.69/chat_title_more_dots.png +0 -0
- screenshot_vision_algorithm/android/wechat/templates/android/8.0.69/favorite_label.png +0 -0
- screenshot_vision_algorithm/android/wechat/templates/android/8.0.69/new_messages_hint_suffix.png +0 -0
- screenshot_vision_algorithm/android/wechat/templates/android/8.0.69/unread_divider_hint.png +0 -0
- screenshot_vision_algorithm/android/wechat/templates/android/8.0.69/unread_divider_hint_v2_textonly.png +0 -0
- screenshot_vision_algorithm/android/wechat/templates/android/8.0.69/wechat_note_header.png +0 -0
- screenshot_vision_algorithm/android/xhs/__init__.py +4 -0
- screenshot_vision_algorithm/android/zhihu/__init__.py +4 -0
- screenshot_vision_algorithm/png_utils.py +86 -0
- screenshot_vision_algorithm-0.3.0.dist-info/METADATA +425 -0
- screenshot_vision_algorithm-0.3.0.dist-info/RECORD +40 -0
- screenshot_vision_algorithm-0.3.0.dist-info/WHEEL +5 -0
- screenshot_vision_algorithm-0.3.0.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,48 @@
|
|
|
1
|
+
"""Public API for screenshot_vision_algorithm.android.wechat."""
|
|
2
|
+
|
|
3
|
+
from screenshot_vision_algorithm._config import Platform, WeChatVersion, Profile
|
|
4
|
+
|
|
5
|
+
from screenshot_vision_algorithm.android.wechat.algorithms.card_bbox import (
|
|
6
|
+
TAP_ANTI_DETECT_X_JITTER_PX,
|
|
7
|
+
_compute_exact_card_and_bubble_bboxes,
|
|
8
|
+
_compute_exact_card_bboxes,
|
|
9
|
+
BubbleBbox,
|
|
10
|
+
FavoriteLabelHit,
|
|
11
|
+
ThumbnailCard,
|
|
12
|
+
TrackedCard,
|
|
13
|
+
derive_cards,
|
|
14
|
+
derive_cards_and_bubbles,
|
|
15
|
+
drop_top_clamped_false_positive_cards,
|
|
16
|
+
bbox_to_metadata_list,
|
|
17
|
+
bubble_to_metadata_list,
|
|
18
|
+
card_bounding_tuple,
|
|
19
|
+
card_overlaps_processed,
|
|
20
|
+
click_context_for_tap_thumbnail,
|
|
21
|
+
pick_first_unprocessed_card,
|
|
22
|
+
pick_next_unclicked_card,
|
|
23
|
+
y_interval_overlap_ratio,
|
|
24
|
+
)
|
|
25
|
+
|
|
26
|
+
__all__ = [
|
|
27
|
+
"Platform",
|
|
28
|
+
"WeChatVersion",
|
|
29
|
+
"Profile",
|
|
30
|
+
"ThumbnailCard",
|
|
31
|
+
"TrackedCard",
|
|
32
|
+
"FavoriteLabelHit",
|
|
33
|
+
"BubbleBbox",
|
|
34
|
+
"TAP_ANTI_DETECT_X_JITTER_PX",
|
|
35
|
+
"_compute_exact_card_and_bubble_bboxes",
|
|
36
|
+
"_compute_exact_card_bboxes",
|
|
37
|
+
"derive_cards",
|
|
38
|
+
"derive_cards_and_bubbles",
|
|
39
|
+
"drop_top_clamped_false_positive_cards",
|
|
40
|
+
"bbox_to_metadata_list",
|
|
41
|
+
"bubble_to_metadata_list",
|
|
42
|
+
"card_bounding_tuple",
|
|
43
|
+
"card_overlaps_processed",
|
|
44
|
+
"click_context_for_tap_thumbnail",
|
|
45
|
+
"pick_first_unprocessed_card",
|
|
46
|
+
"pick_next_unclicked_card",
|
|
47
|
+
"y_interval_overlap_ratio",
|
|
48
|
+
]
|
|
@@ -0,0 +1,61 @@
|
|
|
1
|
+
"""Platform profile: parameter resolution by (platform, wechat_version)."""
|
|
2
|
+
|
|
3
|
+
from dataclasses import dataclass
|
|
4
|
+
from enum import Enum
|
|
5
|
+
from importlib import import_module
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
from typing import Any
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class Platform(str, Enum):
|
|
11
|
+
ANDROID = "android"
|
|
12
|
+
IOS = "ios"
|
|
13
|
+
HARMONY = "harmony"
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class WeChatVersion(str, Enum):
|
|
17
|
+
V8_0_69 = "8.0.69"
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
#: All design constants are anchored to 1080x screen width (WeChat baseline).
|
|
21
|
+
BASELINE_WIDTH = 1080
|
|
22
|
+
|
|
23
|
+
_PROFILE_MODULES = {
|
|
24
|
+
Platform.ANDROID: "screenshot_vision_algorithm.android.wechat.profiles.android",
|
|
25
|
+
Platform.IOS: "screenshot_vision_algorithm.android.wechat.profiles.ios",
|
|
26
|
+
Platform.HARMONY: "screenshot_vision_algorithm.android.wechat.profiles.harmony",
|
|
27
|
+
}
|
|
28
|
+
|
|
29
|
+
_PACKAGE_ROOT = Path(__file__).resolve().parent
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
@dataclass(frozen=True)
|
|
33
|
+
class Profile:
|
|
34
|
+
"""Parameter bundle for one (platform, wechat_version) combination.
|
|
35
|
+
|
|
36
|
+
Usage:
|
|
37
|
+
p = Profile(platform=Platform.ANDROID, wechat_version=WeChatVersion.V8_0_69)
|
|
38
|
+
templates = p.templates_dir # Path to versioned PNG assets
|
|
39
|
+
"""
|
|
40
|
+
|
|
41
|
+
platform: Platform
|
|
42
|
+
wechat_version: WeChatVersion
|
|
43
|
+
|
|
44
|
+
@property
|
|
45
|
+
def _profile_module(self) -> Any:
|
|
46
|
+
name = _PROFILE_MODULES[self.platform]
|
|
47
|
+
return import_module(name)
|
|
48
|
+
|
|
49
|
+
@property
|
|
50
|
+
def templates_dir(self) -> Path:
|
|
51
|
+
return _PACKAGE_ROOT / "android" / "wechat" / "templates" / self.platform.value / self.wechat_version.value
|
|
52
|
+
|
|
53
|
+
def __getattr__(self, name: str) -> Any:
|
|
54
|
+
"""Delegate unknown attributes to the platform profile module."""
|
|
55
|
+
try:
|
|
56
|
+
return getattr(self._profile_module, name)
|
|
57
|
+
except AttributeError:
|
|
58
|
+
raise AttributeError(
|
|
59
|
+
f"Profile {self.platform.value}/{self.wechat_version.value} "
|
|
60
|
+
f"has no attribute '{name}'"
|
|
61
|
+
) from None
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Android platform screenshot vision algorithms."""
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""WeChat screenshot vision algorithms."""
|
|
File without changes
|
|
@@ -0,0 +1,209 @@
|
|
|
1
|
+
"""左列头像圆心检测(列表页 / 群聊页共用几何先验)。
|
|
2
|
+
|
|
3
|
+
列表会话行与群聊发言人归因共用同一套 Hough + median_x 过滤 + y 向去重,
|
|
4
|
+
产出 ``AvatarCentroid`` 供 ROI 守门与 PRD §5 纵向筒使用。
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import math
|
|
10
|
+
from dataclasses import dataclass
|
|
11
|
+
from typing import Optional
|
|
12
|
+
|
|
13
|
+
import numpy as np
|
|
14
|
+
from loguru import logger
|
|
15
|
+
|
|
16
|
+
try:
|
|
17
|
+
import cv2
|
|
18
|
+
except ImportError:
|
|
19
|
+
cv2 = None # type: ignore[assignment]
|
|
20
|
+
|
|
21
|
+
from screenshot_vision_algorithm.android.wechat.algorithms.template_matching import (
|
|
22
|
+
CONVERSATION_LIST_AVATAR_MEDIAN_X_HALF_WIDTH_BASELINE,
|
|
23
|
+
CONVERSATION_LIST_AVATAR_ROI_LEFT_BASELINE,
|
|
24
|
+
CONVERSATION_LIST_AVATAR_ROI_RIGHT_BASELINE,
|
|
25
|
+
)
|
|
26
|
+
|
|
27
|
+
# 与列表页 Hough 对齐(1080 基准);群聊左列头像物理尺寸与列表一致。
|
|
28
|
+
_LIST_HOUGH_MIN_DIST_BASELINE = 117
|
|
29
|
+
_LIST_HOUGH_MIN_R_BASELINE = 36
|
|
30
|
+
_LIST_HOUGH_MAX_R_BASELINE = 80
|
|
31
|
+
_LIST_YMIN_GAP_BASELINE = 118
|
|
32
|
+
_LIST_SAME_ICON_DY_MAX_BASELINE = 18
|
|
33
|
+
|
|
34
|
+
# 绑定:昵称行与头像锚点的最大竖直间距(px @ 1080)。
|
|
35
|
+
NICKNAME_AVATAR_BIND_MAX_DY_BASELINE = 140
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
@dataclass(frozen=True)
|
|
39
|
+
class AvatarCentroid:
|
|
40
|
+
"""左列头像圆(整图像素坐标)。"""
|
|
41
|
+
|
|
42
|
+
cx: int
|
|
43
|
+
cy: int
|
|
44
|
+
r: int
|
|
45
|
+
|
|
46
|
+
@property
|
|
47
|
+
def y_top(self) -> int:
|
|
48
|
+
return int(self.cy - self.r)
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
@dataclass(frozen=True)
|
|
52
|
+
class LeftAvatarColumnLayout:
|
|
53
|
+
centroids: tuple[AvatarCentroid, ...]
|
|
54
|
+
median_cx: float
|
|
55
|
+
column_x_half_width_px: int
|
|
56
|
+
|
|
57
|
+
@property
|
|
58
|
+
def empty(self) -> bool:
|
|
59
|
+
return not self.centroids
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def _sw(scale_w: float) -> float:
|
|
63
|
+
return max(float(scale_w), 1e-6)
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def detect_left_avatar_column_layout(
|
|
67
|
+
screen_bgr: np.ndarray,
|
|
68
|
+
scale_w: float,
|
|
69
|
+
*,
|
|
70
|
+
y_top: int = 0,
|
|
71
|
+
y_bottom_excl: Optional[int] = None,
|
|
72
|
+
) -> LeftAvatarColumnLayout:
|
|
73
|
+
"""在左列 ROI 内 Hough 圆 → median_x 过滤 → y 向去重(与列表页同参)。"""
|
|
74
|
+
sw = _sw(scale_w)
|
|
75
|
+
xt = int(round(float(CONVERSATION_LIST_AVATAR_MEDIAN_X_HALF_WIDTH_BASELINE) * sw))
|
|
76
|
+
empty = LeftAvatarColumnLayout(centroids=(), median_cx=0.0, column_x_half_width_px=xt)
|
|
77
|
+
|
|
78
|
+
if cv2 is None or screen_bgr is None or getattr(screen_bgr, "size", 0) == 0:
|
|
79
|
+
return empty
|
|
80
|
+
|
|
81
|
+
h_full, w_full = screen_bgr.shape[:2]
|
|
82
|
+
y1 = int(max(0, min(y_top, h_full - 1)))
|
|
83
|
+
y2 = int(y_bottom_excl if y_bottom_excl is not None else h_full)
|
|
84
|
+
y2 = int(max(y1 + 8, min(y2, h_full)))
|
|
85
|
+
x1_roi = int(max(0, round(CONVERSATION_LIST_AVATAR_ROI_LEFT_BASELINE * sw)))
|
|
86
|
+
x2_roi = int(min(w_full, round(CONVERSATION_LIST_AVATAR_ROI_RIGHT_BASELINE * sw)))
|
|
87
|
+
if y2 <= y1 + 40 or x2_roi <= x1_roi + 24:
|
|
88
|
+
return empty
|
|
89
|
+
|
|
90
|
+
crop = screen_bgr[y1:y2, x1_roi:x2_roi]
|
|
91
|
+
gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
|
|
92
|
+
gray_blur = cv2.GaussianBlur(gray, (5, 5), 0)
|
|
93
|
+
circles = cv2.HoughCircles(
|
|
94
|
+
gray_blur,
|
|
95
|
+
cv2.HOUGH_GRADIENT,
|
|
96
|
+
dp=1.25,
|
|
97
|
+
minDist=int(round(_LIST_HOUGH_MIN_DIST_BASELINE * sw)),
|
|
98
|
+
param1=100,
|
|
99
|
+
param2=22,
|
|
100
|
+
minRadius=max(8, int(round(_LIST_HOUGH_MIN_R_BASELINE * sw))),
|
|
101
|
+
maxRadius=int(round(_LIST_HOUGH_MAX_R_BASELINE * sw)),
|
|
102
|
+
)
|
|
103
|
+
if circles is None:
|
|
104
|
+
return empty
|
|
105
|
+
|
|
106
|
+
cand: list[tuple[float, float, float]] = []
|
|
107
|
+
for xi, yi, ri in circles[0]:
|
|
108
|
+
cand.append((float(ri), float(xi), float(yi)))
|
|
109
|
+
|
|
110
|
+
if not cand:
|
|
111
|
+
return empty
|
|
112
|
+
|
|
113
|
+
xs_crop = [c[1] for c in cand]
|
|
114
|
+
median_x_crop = float(np.median(np.array(xs_crop, dtype=np.float64)))
|
|
115
|
+
median_cx_full = float(x1_roi) + median_x_crop
|
|
116
|
+
filt = [c for c in cand if abs(c[1] - median_x_crop) <= xt]
|
|
117
|
+
if not filt:
|
|
118
|
+
return empty
|
|
119
|
+
|
|
120
|
+
filt.sort(key=lambda t: t[2])
|
|
121
|
+
ymin_gap = round(_LIST_YMIN_GAP_BASELINE * sw)
|
|
122
|
+
same_dy = float(_LIST_SAME_ICON_DY_MAX_BASELINE) * sw
|
|
123
|
+
stacked: list[tuple[float, float, float]] = []
|
|
124
|
+
for r, xc, yc in filt:
|
|
125
|
+
if not stacked:
|
|
126
|
+
stacked.append((r, xc, yc))
|
|
127
|
+
continue
|
|
128
|
+
pr, _px, py_prev = stacked[-1]
|
|
129
|
+
dy = float(yc) - float(py_prev)
|
|
130
|
+
if dy < ymin_gap:
|
|
131
|
+
if dy <= same_dy and r > pr:
|
|
132
|
+
stacked[-1] = (r, xc, yc)
|
|
133
|
+
else:
|
|
134
|
+
stacked.append((r, xc, yc))
|
|
135
|
+
|
|
136
|
+
centroids: list[AvatarCentroid] = []
|
|
137
|
+
for r, xc, yc in stacked:
|
|
138
|
+
cx_full = int(round(x1_roi + xc))
|
|
139
|
+
cy_full = int(round(y1 + yc))
|
|
140
|
+
ri = int(round(r))
|
|
141
|
+
centroids.append(AvatarCentroid(cx=cx_full, cy=cy_full, r=max(8, ri)))
|
|
142
|
+
|
|
143
|
+
return LeftAvatarColumnLayout(
|
|
144
|
+
centroids=tuple(centroids),
|
|
145
|
+
median_cx=median_cx_full,
|
|
146
|
+
column_x_half_width_px=xt,
|
|
147
|
+
)
|
|
148
|
+
|
|
149
|
+
|
|
150
|
+
def find_avatar_anchor_for_nickname_bbox(
|
|
151
|
+
layout: LeftAvatarColumnLayout,
|
|
152
|
+
bbox_xyxy: tuple[int, int, int, int],
|
|
153
|
+
*,
|
|
154
|
+
scale_w: float,
|
|
155
|
+
max_dy_px: Optional[int] = None,
|
|
156
|
+
) -> Optional[AvatarCentroid]:
|
|
157
|
+
"""为昵称行找最近且在其上方的左列头像(几何绑定)。"""
|
|
158
|
+
if layout.empty:
|
|
159
|
+
return None
|
|
160
|
+
x1, y1, x2, y2 = bbox_xyxy
|
|
161
|
+
nick_cy = (y1 + y2) / 2.0
|
|
162
|
+
sw = _sw(scale_w)
|
|
163
|
+
max_dy = (
|
|
164
|
+
int(max_dy_px)
|
|
165
|
+
if max_dy_px is not None
|
|
166
|
+
else math.floor(NICKNAME_AVATAR_BIND_MAX_DY_BASELINE * sw)
|
|
167
|
+
)
|
|
168
|
+
nick_w = max(1, x2 - x1)
|
|
169
|
+
x_hi = float(layout.median_cx) + float(layout.column_x_half_width_px) + nick_w * 2
|
|
170
|
+
|
|
171
|
+
best: Optional[AvatarCentroid] = None
|
|
172
|
+
best_score = 1e18
|
|
173
|
+
|
|
174
|
+
for c in layout.centroids:
|
|
175
|
+
if float(c.cx) > x_hi:
|
|
176
|
+
continue
|
|
177
|
+
dy = float(y1) - float(c.cy)
|
|
178
|
+
if dy < -100:
|
|
179
|
+
continue
|
|
180
|
+
if dy >= max_dy:
|
|
181
|
+
continue
|
|
182
|
+
score = abs(float(nick_cy) - float(c.cy))
|
|
183
|
+
if score < best_score:
|
|
184
|
+
best_score = score
|
|
185
|
+
best = c
|
|
186
|
+
|
|
187
|
+
return best
|
|
188
|
+
|
|
189
|
+
|
|
190
|
+
def nickname_bbox_in_avatar_column(
|
|
191
|
+
layout: LeftAvatarColumnLayout,
|
|
192
|
+
bbox_xyxy: tuple[int, int, int, int],
|
|
193
|
+
*,
|
|
194
|
+
original_width: int,
|
|
195
|
+
max_x1_ratio: float = 0.30,
|
|
196
|
+
) -> bool:
|
|
197
|
+
"""昵称行是否落在左列头像带(median_x 优先,无检出时回退 x1 比例)。"""
|
|
198
|
+
x1, _, x2, _ = bbox_xyxy
|
|
199
|
+
if layout.empty or layout.median_cx <= 0:
|
|
200
|
+
return float(x1) < float(original_width) * float(max_x1_ratio)
|
|
201
|
+
nick_w = max(1, x2 - x1)
|
|
202
|
+
lo = float(layout.median_cx) - float(layout.column_x_half_width_px)
|
|
203
|
+
hi = float(layout.median_cx) + float(layout.column_x_half_width_px) + nick_w * 2
|
|
204
|
+
return lo <= float(x1) <= hi
|
|
205
|
+
|
|
206
|
+
|
|
207
|
+
def left_avatar_ytops(layout: LeftAvatarColumnLayout) -> list[int]:
|
|
208
|
+
"""供 PRD §5:头像顶 y(圆心减半径)。"""
|
|
209
|
+
return [c.y_top for c in layout.centroids]
|
|
@@ -0,0 +1,275 @@
|
|
|
1
|
+
"""Experimental list unread-badge heuristic: horizontal band × HSV red × bright glyphs.
|
|
2
|
+
|
|
3
|
+
Used for calibration / offline scripts (not wired into scan_pinned_groups by default).
|
|
4
|
+
Aligned with UX: badge sits on avatar left column — constrain x as fractions of screen width.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import math
|
|
10
|
+
from dataclasses import dataclass
|
|
11
|
+
|
|
12
|
+
import cv2
|
|
13
|
+
import numpy as np
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
@dataclass(frozen=True)
|
|
17
|
+
class WhiteOnRedBadgeCandidate:
|
|
18
|
+
"""One candidate unread badge (~= disk-like red blob with white glyphs)."""
|
|
19
|
+
|
|
20
|
+
x: int
|
|
21
|
+
y: int
|
|
22
|
+
w: int
|
|
23
|
+
h: int
|
|
24
|
+
contour_area: int
|
|
25
|
+
bright_ratio: float
|
|
26
|
+
#: ``4*pi*area/perimeter²`` (~1 disk, ~0.78 square)
|
|
27
|
+
circularity: float
|
|
28
|
+
#: ``max(fitEllipse axes) / min(axes)`` (≈1 circle; >>1 elongated)
|
|
29
|
+
ellipse_axis_ratio: float
|
|
30
|
+
score: float
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def _dual_red_mask(hsv_roi: np.ndarray, *, sat_min: int, val_min: int, low_h_max: int) -> np.ndarray:
|
|
34
|
+
low_lo = np.array((0, sat_min, val_min), dtype=np.uint8)
|
|
35
|
+
low_hi = np.array((low_h_max, 255, 255), dtype=np.uint8)
|
|
36
|
+
high_lo = np.array((170, sat_min, val_min), dtype=np.uint8)
|
|
37
|
+
high_hi = np.array((180, 255, 255), dtype=np.uint8)
|
|
38
|
+
return cv2.bitwise_or(
|
|
39
|
+
cv2.inRange(hsv_roi, low_lo, low_hi),
|
|
40
|
+
cv2.inRange(hsv_roi, high_lo, high_hi),
|
|
41
|
+
)
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def _contour_circularity_pixels(contour: np.ndarray) -> float:
|
|
45
|
+
"""Isoperimetric quotient; ~1 full disk, lower for dented/skinny blobs."""
|
|
46
|
+
area = abs(float(cv2.contourArea(contour)))
|
|
47
|
+
peri = float(cv2.arcLength(contour, True))
|
|
48
|
+
if peri < 1e-6:
|
|
49
|
+
return 0.0
|
|
50
|
+
return float((4.0 * math.pi * max(area, 1e-9)) / (peri * peri))
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
def _ellipse_axis_ratio_or_none(contour: np.ndarray) -> float | None:
|
|
54
|
+
"""``max(fitEllipse axis) / min(axis)``. ``None`` if ``len < 5``."""
|
|
55
|
+
if len(contour) < 5:
|
|
56
|
+
return None
|
|
57
|
+
(_, _), (maj, mn), _ = cv2.fitEllipse(contour)
|
|
58
|
+
mj, nn = float(maj), float(mn)
|
|
59
|
+
if mj < 1e-9 or nn < 1e-9:
|
|
60
|
+
return None
|
|
61
|
+
return float(max(mj, nn) / min(mj, nn))
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
def count_badges_white_on_red(
|
|
65
|
+
screen_bgr: np.ndarray,
|
|
66
|
+
*,
|
|
67
|
+
width_frac_lo: float = 0.085,
|
|
68
|
+
width_frac_hi: float = 0.195,
|
|
69
|
+
vertical_y1: int = 270,
|
|
70
|
+
vertical_y2: int | None = None,
|
|
71
|
+
sat_min: int = 88,
|
|
72
|
+
val_min: int = 82,
|
|
73
|
+
low_h_max: int = 15,
|
|
74
|
+
min_area_px: int = 150,
|
|
75
|
+
max_area_px: int = 4000,
|
|
76
|
+
max_aspect: float = 1.95,
|
|
77
|
+
min_bright_ratio: float = 0.04,
|
|
78
|
+
white_gray_floor: int = 204,
|
|
79
|
+
morph_extra_dilate_iters: int = 1,
|
|
80
|
+
max_bbox_center_x_frac: float | None = 0.185,
|
|
81
|
+
min_bbox_center_y_abs: int | None = None,
|
|
82
|
+
max_bright_ratio: float | None = 0.17,
|
|
83
|
+
) -> tuple[tuple[WhiteOnRedBadgeCandidate, ...], tuple[int, int, int, int]]:
|
|
84
|
+
"""Find unread-style badges inside an x-percent band.
|
|
85
|
+
|
|
86
|
+
Steps (per user's scheme):
|
|
87
|
+
1) Restrict to horizontal band ``[width_frac_lo, width_frac_hi]`` of screen width
|
|
88
|
+
plus optional vertical clipping (conversation list bodies).
|
|
89
|
+
2) Dual-band HSV red mask + morphology.
|
|
90
|
+
3) Contour filter (area / aspect).
|
|
91
|
+
4) Reject blobs where the **filled contour** lacks enough luminous pixels
|
|
92
|
+
(**bright_ratio** = #(gray≥``white_gray_floor``∩red fill) / #red fill pixels).
|
|
93
|
+
|
|
94
|
+
Returns:
|
|
95
|
+
Sorted candidates top→bottom, and ``(roi_x1, roi_y1, roi_x2, roi_y2)``.
|
|
96
|
+
"""
|
|
97
|
+
h_full, w_full = screen_bgr.shape[:2]
|
|
98
|
+
vx2 = vertical_y2 if vertical_y2 is not None else h_full - 120
|
|
99
|
+
|
|
100
|
+
roi_x1 = max(0, int(w_full * width_frac_lo))
|
|
101
|
+
roi_x2 = min(w_full, int(w_full * width_frac_hi))
|
|
102
|
+
roi_y1 = max(0, vertical_y1)
|
|
103
|
+
roi_y2 = min(h_full, vx2)
|
|
104
|
+
|
|
105
|
+
if roi_x2 <= roi_x1 or roi_y2 <= roi_y1:
|
|
106
|
+
return (), (0, 0, 0, 0)
|
|
107
|
+
|
|
108
|
+
roi_rgb = screen_bgr[roi_y1:roi_y2, roi_x1:roi_x2]
|
|
109
|
+
gray_roi = cv2.cvtColor(roi_rgb, cv2.COLOR_BGR2GRAY)
|
|
110
|
+
|
|
111
|
+
hsv = cv2.cvtColor(roi_rgb, cv2.COLOR_BGR2HSV)
|
|
112
|
+
mask = _dual_red_mask(hsv, sat_min=sat_min, val_min=val_min, low_h_max=low_h_max)
|
|
113
|
+
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
|
|
114
|
+
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
|
|
115
|
+
if morph_extra_dilate_iters > 0:
|
|
116
|
+
mask = cv2.dilate(mask, kernel, iterations=morph_extra_dilate_iters)
|
|
117
|
+
|
|
118
|
+
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
119
|
+
hits: list[WhiteOnRedBadgeCandidate] = []
|
|
120
|
+
|
|
121
|
+
for c in contours:
|
|
122
|
+
rx, ry, rw, rh = cv2.boundingRect(c)
|
|
123
|
+
area = float(cv2.contourArea(c))
|
|
124
|
+
if area < min_area_px or area > max_area_px:
|
|
125
|
+
continue
|
|
126
|
+
ls, ss = max(rw, rh), max(1, min(rw, rh))
|
|
127
|
+
if ls / ss > max_aspect:
|
|
128
|
+
continue
|
|
129
|
+
|
|
130
|
+
cx = rx + rw // 2
|
|
131
|
+
cy = ry + rh // 2
|
|
132
|
+
if max_bbox_center_x_frac is not None:
|
|
133
|
+
if (roi_x1 + cx) > int(w_full * max_bbox_center_x_frac):
|
|
134
|
+
continue
|
|
135
|
+
gc_y_abs = roi_y1 + cy
|
|
136
|
+
if min_bbox_center_y_abs is not None and gc_y_abs < min_bbox_center_y_abs:
|
|
137
|
+
continue
|
|
138
|
+
|
|
139
|
+
# White-ish digits: luminous pixels strictly inside filled red silhouette.
|
|
140
|
+
h_roi0, w_roi0 = gray_roi.shape
|
|
141
|
+
fill = np.zeros((h_roi0, w_roi0), dtype=np.uint8)
|
|
142
|
+
cv2.drawContours(fill, [c], 0, 255, -1)
|
|
143
|
+
bright_on_fill = np.count_nonzero(
|
|
144
|
+
np.logical_and(fill > 0, gray_roi >= np.uint8(white_gray_floor))
|
|
145
|
+
)
|
|
146
|
+
red_px = max(1, int(cv2.countNonZero(fill)))
|
|
147
|
+
bright_ratio = float(bright_on_fill) / float(red_px)
|
|
148
|
+
|
|
149
|
+
if bright_ratio < min_bright_ratio:
|
|
150
|
+
continue
|
|
151
|
+
if max_bright_ratio is not None and bright_ratio > max_bright_ratio:
|
|
152
|
+
continue
|
|
153
|
+
|
|
154
|
+
circ = _contour_circularity_pixels(c)
|
|
155
|
+
ear_raw = _ellipse_axis_ratio_or_none(c)
|
|
156
|
+
ear = float(ear_raw) if ear_raw is not None else 1.0
|
|
157
|
+
|
|
158
|
+
score = area * bright_ratio
|
|
159
|
+
|
|
160
|
+
gx = roi_x1 + rx
|
|
161
|
+
gy = roi_y1 + ry
|
|
162
|
+
hits.append(
|
|
163
|
+
WhiteOnRedBadgeCandidate(
|
|
164
|
+
x=int(gx),
|
|
165
|
+
y=int(gy),
|
|
166
|
+
w=int(rw),
|
|
167
|
+
h=int(rh),
|
|
168
|
+
contour_area=int(area),
|
|
169
|
+
bright_ratio=float(bright_ratio),
|
|
170
|
+
circularity=float(circ),
|
|
171
|
+
ellipse_axis_ratio=ear,
|
|
172
|
+
score=float(score),
|
|
173
|
+
)
|
|
174
|
+
)
|
|
175
|
+
|
|
176
|
+
hits.sort(key=lambda u: u.y)
|
|
177
|
+
merged = _dedupe_by_vertical_bucket(hits, bucket_px=95)
|
|
178
|
+
return tuple(merged), (roi_x1, roi_y1, roi_x2, roi_y2)
|
|
179
|
+
|
|
180
|
+
|
|
181
|
+
def _dedupe_by_vertical_bucket(
|
|
182
|
+
hits: list[WhiteOnRedBadgeCandidate],
|
|
183
|
+
*,
|
|
184
|
+
bucket_px: float,
|
|
185
|
+
) -> list[WhiteOnRedBadgeCandidate]:
|
|
186
|
+
"""Merge multiple contours from the same list row (~one badge split)."""
|
|
187
|
+
|
|
188
|
+
if not hits:
|
|
189
|
+
return []
|
|
190
|
+
|
|
191
|
+
buckets: dict[int, list[WhiteOnRedBadgeCandidate]] = {}
|
|
192
|
+
for h in hits:
|
|
193
|
+
key = int((h.y + h.h // 2) // bucket_px)
|
|
194
|
+
buckets.setdefault(key, []).append(h)
|
|
195
|
+
|
|
196
|
+
merged: list[WhiteOnRedBadgeCandidate] = []
|
|
197
|
+
for grp in buckets.values():
|
|
198
|
+
grp.sort(key=lambda u: (-u.score, -u.contour_area))
|
|
199
|
+
merged.append(grp[0])
|
|
200
|
+
merged.sort(key=lambda u: u.y)
|
|
201
|
+
return merged
|
|
202
|
+
|
|
203
|
+
|
|
204
|
+
def count_badges_per_row_roi(
|
|
205
|
+
screen_bgr: np.ndarray,
|
|
206
|
+
*,
|
|
207
|
+
scale_w: float | None = None,
|
|
208
|
+
min_bright_ratio: float = 0.04,
|
|
209
|
+
min_area_px: int = 260,
|
|
210
|
+
max_bright_ratio: float | None = 0.17,
|
|
211
|
+
**kw: object,
|
|
212
|
+
) -> tuple[tuple[WhiteOnRedBadgeCandidate, ...], tuple[int, int, int, int]]:
|
|
213
|
+
"""Same heuristic but matches :meth:`scan_pinned_groups` per-row avatar band.
|
|
214
|
+
|
|
215
|
+
Reuses DD row-split + ``badge_roi`` horizontal strip (stronger than one
|
|
216
|
+
flat width-% box).
|
|
217
|
+
"""
|
|
218
|
+
|
|
219
|
+
del kw # reserved for forwards-compat
|
|
220
|
+
|
|
221
|
+
from collector.driver import (
|
|
222
|
+
AVATAR_UNREAD_BADGE_LEFT_INSET_BASELINE,
|
|
223
|
+
AVATAR_UNREAD_BADGE_RIGHT_PAD_BASELINE,
|
|
224
|
+
CONVERSATION_AVATAR_RIGHT_EDGE_BASELINE,
|
|
225
|
+
)
|
|
226
|
+
from screenshot_vision_algorithm.android.wechat.algorithms.template_matching import split_conversation_rows
|
|
227
|
+
|
|
228
|
+
h_full, w_full = screen_bgr.shape[:2]
|
|
229
|
+
sw = scale_w if scale_w is not None else float(w_full) / 1080.0
|
|
230
|
+
row_rects = split_conversation_rows(h_full, sw)
|
|
231
|
+
avatar_edge = int(round(CONVERSATION_AVATAR_RIGHT_EDGE_BASELINE * sw))
|
|
232
|
+
ain = int(round(AVATAR_UNREAD_BADGE_LEFT_INSET_BASELINE * sw))
|
|
233
|
+
arp = int(round(AVATAR_UNREAD_BADGE_RIGHT_PAD_BASELINE * sw))
|
|
234
|
+
|
|
235
|
+
all_hits: list[WhiteOnRedBadgeCandidate] = []
|
|
236
|
+
for row_idx, (y1, y2) in enumerate(row_rects):
|
|
237
|
+
rh = y2 - y1
|
|
238
|
+
vy1 = y1 + max(1, int(rh * 0.04))
|
|
239
|
+
vy2 = y1 + int(rh * 0.58)
|
|
240
|
+
vy2 = max(vy1 + 8, vy2)
|
|
241
|
+
x1_bb = max(0, avatar_edge - ain)
|
|
242
|
+
x2_bb = min(w_full, avatar_edge + arp)
|
|
243
|
+
frac_lo = float(x1_bb) / float(w_full)
|
|
244
|
+
frac_hi = float(x2_bb) / float(w_full)
|
|
245
|
+
# Below the toolbar "微信(N)" reds (edb1a89f @1080-wide: ~282–294)
|
|
246
|
+
row_min_cy_abs: int | None = y1 + 18 if row_idx == 0 else None
|
|
247
|
+
cand, roi = count_badges_white_on_red(
|
|
248
|
+
screen_bgr,
|
|
249
|
+
width_frac_lo=frac_lo,
|
|
250
|
+
width_frac_hi=frac_hi,
|
|
251
|
+
vertical_y1=vy1,
|
|
252
|
+
vertical_y2=vy2,
|
|
253
|
+
min_bright_ratio=min_bright_ratio,
|
|
254
|
+
min_area_px=min_area_px,
|
|
255
|
+
max_bright_ratio=max_bright_ratio,
|
|
256
|
+
min_bbox_center_y_abs=row_min_cy_abs,
|
|
257
|
+
)
|
|
258
|
+
if cand:
|
|
259
|
+
all_hits.append(cand[0])
|
|
260
|
+
|
|
261
|
+
x_lo = max(0, int(round((CONVERSATION_AVATAR_RIGHT_EDGE_BASELINE - AVATAR_UNREAD_BADGE_LEFT_INSET_BASELINE) * sw)))
|
|
262
|
+
x_hi = min(
|
|
263
|
+
w_full,
|
|
264
|
+
int(round((CONVERSATION_AVATAR_RIGHT_EDGE_BASELINE + AVATAR_UNREAD_BADGE_RIGHT_PAD_BASELINE) * sw)),
|
|
265
|
+
)
|
|
266
|
+
y_lo = row_rects[0][0] if row_rects else 0
|
|
267
|
+
outer_roi = (x_lo, y_lo, x_hi, h_full)
|
|
268
|
+
return tuple(all_hits), outer_roi
|
|
269
|
+
|
|
270
|
+
|
|
271
|
+
__all__ = [
|
|
272
|
+
"WhiteOnRedBadgeCandidate",
|
|
273
|
+
"count_badges_white_on_red",
|
|
274
|
+
"count_badges_per_row_roi",
|
|
275
|
+
]
|