wechat-screenshot-vision-algorithm 0.1.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.
Files changed (36) hide show
  1. wechat_screenshot_vision_algorithm/__init__.py +40 -0
  2. wechat_screenshot_vision_algorithm/_config.py +61 -0
  3. wechat_screenshot_vision_algorithm/algorithms/__init__.py +0 -0
  4. wechat_screenshot_vision_algorithm/algorithms/avatar_column.py +211 -0
  5. wechat_screenshot_vision_algorithm/algorithms/badge_detection.py +275 -0
  6. wechat_screenshot_vision_algorithm/algorithms/card_bbox.py +814 -0
  7. wechat_screenshot_vision_algorithm/algorithms/phash_utils.py +267 -0
  8. wechat_screenshot_vision_algorithm/algorithms/speaker_band.py +292 -0
  9. wechat_screenshot_vision_algorithm/algorithms/template_matching.py +2152 -0
  10. wechat_screenshot_vision_algorithm/algorithms/title_ocr.py +145 -0
  11. wechat_screenshot_vision_algorithm/merge/__init__.py +0 -0
  12. wechat_screenshot_vision_algorithm/merge/multipage.py +157 -0
  13. wechat_screenshot_vision_algorithm/ocr/__init__.py +0 -0
  14. wechat_screenshot_vision_algorithm/ocr/avatar_guard.py +436 -0
  15. wechat_screenshot_vision_algorithm/ocr/badge_ocr.py +234 -0
  16. wechat_screenshot_vision_algorithm/ocr/nickname_binding.py +1888 -0
  17. wechat_screenshot_vision_algorithm/ocr/text_ocr_adapter.py +627 -0
  18. wechat_screenshot_vision_algorithm/png_utils.py +87 -0
  19. wechat_screenshot_vision_algorithm/profiles/__init__.py +0 -0
  20. wechat_screenshot_vision_algorithm/profiles/android_wechat.py +53 -0
  21. wechat_screenshot_vision_algorithm/profiles/harmony_wechat.py +10 -0
  22. wechat_screenshot_vision_algorithm/profiles/ios_wechat.py +53 -0
  23. wechat_screenshot_vision_algorithm/templates/wechat/android/8.0.69/chat_back_chevron.png +0 -0
  24. wechat_screenshot_vision_algorithm/templates/wechat/android/8.0.69/chat_input_emoji_smile.png +0 -0
  25. wechat_screenshot_vision_algorithm/templates/wechat/android/8.0.69/chat_input_plus.png +0 -0
  26. wechat_screenshot_vision_algorithm/templates/wechat/android/8.0.69/chat_input_voice.png +0 -0
  27. wechat_screenshot_vision_algorithm/templates/wechat/android/8.0.69/chat_title_more_dots.png +0 -0
  28. wechat_screenshot_vision_algorithm/templates/wechat/android/8.0.69/favorite_label.png +0 -0
  29. wechat_screenshot_vision_algorithm/templates/wechat/android/8.0.69/new_messages_hint_suffix.png +0 -0
  30. wechat_screenshot_vision_algorithm/templates/wechat/android/8.0.69/unread_divider_hint.png +0 -0
  31. wechat_screenshot_vision_algorithm/templates/wechat/android/8.0.69/unread_divider_hint_v2_textonly.png +0 -0
  32. wechat_screenshot_vision_algorithm/templates/wechat/android/8.0.69/wechat_note_header.png +0 -0
  33. wechat_screenshot_vision_algorithm-0.1.0.dist-info/METADATA +423 -0
  34. wechat_screenshot_vision_algorithm-0.1.0.dist-info/RECORD +36 -0
  35. wechat_screenshot_vision_algorithm-0.1.0.dist-info/WHEEL +5 -0
  36. wechat_screenshot_vision_algorithm-0.1.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,436 @@
1
+ """昵称行左侧「头像栏」影像守门:标准 ROI + 可选 Hough 头像锚点对齐 ROI。
2
+
3
+ 纹理(Laplacian + RGB)与列边缘形状须 **同时** 满足(布尔过滤,非算分)。
4
+ """
5
+
6
+ from __future__ import annotations
7
+
8
+ import logging
9
+ from typing import Any, Literal, Optional
10
+
11
+ Side = Literal["left", "right"]
12
+
13
+ import numpy as np
14
+
15
+ logger = logging.getLogger("wx_processor.nickname_avatar_guard")
16
+
17
+ try:
18
+ import cv2
19
+ except ImportError:
20
+ cv2 = None # type: ignore[assignment]
21
+
22
+ try:
23
+ from processor.left_avatar_column import AvatarCentroid
24
+ except ImportError:
25
+ AvatarCentroid = None # type: ignore[misc, assignment]
26
+
27
+
28
+ def bgr_roi_laplacian_variance(roi: np.ndarray) -> float:
29
+ if cv2 is None or roi.size == 0 or roi.shape[0] < 6 or roi.shape[1] < 6:
30
+ return 0.0
31
+ try:
32
+ roi = np.ascontiguousarray(roi, dtype=np.uint8)
33
+ gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
34
+ blur = cv2.GaussianBlur(gray, (3, 3), 0)
35
+ return float(cv2.Laplacian(blur, cv2.CV_64F).var())
36
+ except Exception:
37
+ return 0.0
38
+
39
+
40
+ def bgr_roi_color_std_mean(roi: np.ndarray) -> float:
41
+ flat = roi.reshape(-1, 3).astype(np.float32)
42
+ if flat.size == 0:
43
+ return 0.0
44
+ return float(np.std(flat, axis=0).mean())
45
+
46
+
47
+ def bgr_roi_edge_shape_ok(
48
+ roi: np.ndarray,
49
+ *,
50
+ sig_cols_min: float = 15.0,
51
+ right_left_ratio_min: float = 1.2,
52
+ ) -> bool:
53
+ if cv2 is None or roi.size == 0 or roi.shape[0] < 6 or roi.shape[1] < 6:
54
+ return False
55
+ try:
56
+ roi = np.ascontiguousarray(roi, dtype=np.uint8)
57
+ gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
58
+ sobelx = cv2.Sobel(gray.astype(np.float32), cv2.CV_64F, 1, 0, ksize=3)
59
+ except Exception:
60
+ return False
61
+ col_edge = np.mean(np.abs(sobelx), axis=0)
62
+ if col_edge.size == 0:
63
+ return False
64
+ third = max(1, len(col_edge) // 3)
65
+ left_avg = float(np.mean(col_edge[:third]))
66
+ right_avg = float(np.mean(col_edge[2 * third :]))
67
+ edge_thresh = max(float(np.percentile(col_edge, 60)), 3.0)
68
+ sig_cols = float(np.sum(col_edge > edge_thresh))
69
+ right_left_ratio = right_avg / max(left_avg, 1.0)
70
+ return bool(sig_cols >= sig_cols_min and right_left_ratio > right_left_ratio_min)
71
+
72
+
73
+ def _roi_passes_texture(
74
+ roi: np.ndarray,
75
+ *,
76
+ laplacian_min: float,
77
+ rgb_std_min: float,
78
+ ) -> bool:
79
+ lap_var = bgr_roi_laplacian_variance(roi)
80
+ rgb_std = bgr_roi_color_std_mean(roi)
81
+ return lap_var >= float(laplacian_min) and rgb_std >= float(rgb_std_min)
82
+
83
+
84
+ def _roi_passes_texture_and_shape(
85
+ roi: np.ndarray,
86
+ *,
87
+ laplacian_min: float,
88
+ rgb_std_min: float,
89
+ edge_sig_cols_min: float,
90
+ edge_right_left_ratio_min: float,
91
+ ) -> bool:
92
+ if not _roi_passes_texture(roi, laplacian_min=laplacian_min, rgb_std_min=rgb_std_min):
93
+ return False
94
+ return bgr_roi_edge_shape_ok(
95
+ roi,
96
+ sig_cols_min=edge_sig_cols_min,
97
+ right_left_ratio_min=edge_right_left_ratio_min,
98
+ )
99
+
100
+
101
+ def _crop_standard_roi(
102
+ bgr: np.ndarray,
103
+ *,
104
+ nickname_x1: int,
105
+ y_top: int,
106
+ y_bottom: int,
107
+ gutter_px: int,
108
+ min_roi_width_px: int,
109
+ ) -> Optional[np.ndarray]:
110
+ h_i, w_i = bgr.shape[:2]
111
+ x_end = int(max(0, min(int(nickname_x1) - int(gutter_px), w_i)))
112
+ if x_end < int(min_roi_width_px):
113
+ return None
114
+ y_t = int(max(0, min(y_top, h_i - 1)))
115
+ y_b = int(max(y_t + 4, min(y_bottom, h_i)))
116
+ roi = np.ascontiguousarray(bgr[y_t:y_b, 0:x_end])
117
+ if roi.size == 0 or roi.shape[1] < int(min_roi_width_px):
118
+ return None
119
+ return roi
120
+
121
+
122
+ def _crop_anchor_aligned_roi(
123
+ bgr: np.ndarray,
124
+ *,
125
+ anchor: "AvatarCentroid",
126
+ nickname_x1: int,
127
+ y_bottom: int,
128
+ gutter_px: int,
129
+ min_roi_width_px: int,
130
+ ) -> Optional[np.ndarray]:
131
+ """头像圆顶 → 昵称行底;与列表页「先定位头像再裁窗」一致。"""
132
+ h_i, w_i = bgr.shape[:2]
133
+ x_end = int(max(0, min(int(nickname_x1) - int(gutter_px), w_i)))
134
+ if x_end < int(min_roi_width_px):
135
+ return None
136
+ y_t = int(max(0, min(anchor.y_top, h_i - 1)))
137
+ y_b = int(max(y_t + 4, min(y_bottom, h_i)))
138
+ roi = np.ascontiguousarray(bgr[y_t:y_b, 0:x_end])
139
+ if roi.size == 0 or roi.shape[1] < int(min_roi_width_px):
140
+ return None
141
+ return roi
142
+
143
+
144
+ def _crop_avatar_disk_roi(
145
+ bgr: np.ndarray,
146
+ *,
147
+ anchor: "AvatarCentroid",
148
+ nickname_x1: int,
149
+ gutter_px: int,
150
+ min_roi_width_px: int,
151
+ ) -> Optional[np.ndarray]:
152
+ """头像圆附近紧凑窗,专供形状守门(避免锚点条带过高稀释列边缘)。"""
153
+ h_i, w_i = bgr.shape[:2]
154
+ pad = max(4, int(round(anchor.r * 0.15)))
155
+ x_end = int(max(0, min(int(nickname_x1) - int(gutter_px), w_i)))
156
+ x_start = int(max(0, anchor.cx - anchor.r - pad))
157
+ if x_end - x_start < int(min_roi_width_px):
158
+ return None
159
+ y_t = int(max(0, anchor.cy - anchor.r - pad))
160
+ y_b = int(min(h_i, anchor.cy + anchor.r + pad))
161
+ if y_b <= y_t + 4:
162
+ return None
163
+ roi = np.ascontiguousarray(bgr[y_t:y_b, x_start:x_end])
164
+ if roi.size == 0 or roi.shape[1] < int(min_roi_width_px):
165
+ return None
166
+ return roi
167
+
168
+
169
+ def _crop_standard_roi_right(
170
+ bgr: np.ndarray,
171
+ *,
172
+ nickname_x2: int,
173
+ y_top: int,
174
+ y_bottom: int,
175
+ gutter_px: int,
176
+ min_roi_width_px: int,
177
+ ) -> Optional[np.ndarray]:
178
+ """靠右发言:昵称行右缘以右至屏右缘(与左栏 ``_crop_standard_roi`` 对称)。"""
179
+ h_i, w_i = bgr.shape[:2]
180
+ x_start = int(min(w_i, max(0, int(nickname_x2) + int(gutter_px))))
181
+ if w_i - x_start < int(min_roi_width_px):
182
+ return None
183
+ y_t = int(max(0, min(y_top, h_i - 1)))
184
+ y_b = int(max(y_t + 4, min(y_bottom, h_i)))
185
+ roi = np.ascontiguousarray(bgr[y_t:y_b, x_start:w_i])
186
+ if roi.size == 0 or roi.shape[1] < int(min_roi_width_px):
187
+ return None
188
+ return roi
189
+
190
+
191
+ def nickname_row_passes_prd_avatar_guard(
192
+ bgr: Optional[np.ndarray],
193
+ *,
194
+ nickname_x1: int,
195
+ nickname_x2: int,
196
+ y_top: int,
197
+ y_bottom: int,
198
+ side: Side,
199
+ avatar_anchor: Optional["AvatarCentroid"] = None,
200
+ gutter_px: int = 10,
201
+ min_roi_width_px: int = 40,
202
+ laplacian_min: float = 48.0,
203
+ rgb_std_min: float = 4.0,
204
+ ) -> bool:
205
+ """PRD §九 7(2):**须**有 Hough 头像锚点,且锚点条带 + 圆盘纹理均通过(**禁止**仅靠标准条带纹理∧形状即放行)。"""
206
+ if cv2 is None or bgr is None or bgr.size == 0:
207
+ return False
208
+ if avatar_anchor is None:
209
+ return False
210
+ tex_kw = dict(laplacian_min=laplacian_min, rgb_std_min=rgb_std_min)
211
+ if side == "left":
212
+ roi_strip = _crop_anchor_aligned_roi(
213
+ bgr,
214
+ anchor=avatar_anchor,
215
+ nickname_x1=nickname_x1,
216
+ y_bottom=y_bottom,
217
+ gutter_px=gutter_px,
218
+ min_roi_width_px=min_roi_width_px,
219
+ )
220
+ roi_disk = _crop_avatar_disk_roi(
221
+ bgr,
222
+ anchor=avatar_anchor,
223
+ nickname_x1=nickname_x1,
224
+ gutter_px=gutter_px,
225
+ min_roi_width_px=min_roi_width_px,
226
+ )
227
+ else:
228
+ roi_strip = _crop_anchor_aligned_roi(
229
+ bgr,
230
+ anchor=avatar_anchor,
231
+ nickname_x1=nickname_x2,
232
+ y_bottom=y_bottom,
233
+ gutter_px=gutter_px,
234
+ min_roi_width_px=min_roi_width_px,
235
+ )
236
+ roi_disk = _crop_avatar_disk_roi(
237
+ bgr,
238
+ anchor=avatar_anchor,
239
+ nickname_x1=nickname_x2,
240
+ gutter_px=gutter_px,
241
+ min_roi_width_px=min_roi_width_px,
242
+ )
243
+ if roi_strip is not None and roi_disk is not None:
244
+ if _roi_passes_texture(roi_strip, **tex_kw) and _roi_passes_texture(
245
+ roi_disk, **tex_kw
246
+ ):
247
+ return True
248
+ # 已绑定 Hough 锚点但圆盘纹理偏弱时:仍要求标准条带纹理∧形状(**禁止**无锚点条带放行)。
249
+ shape_kw = dict(
250
+ laplacian_min=laplacian_min,
251
+ rgb_std_min=rgb_std_min,
252
+ edge_sig_cols_min=15.0,
253
+ edge_right_left_ratio_min=1.2,
254
+ )
255
+ if side == "left":
256
+ roi_std = _crop_standard_roi(
257
+ bgr,
258
+ nickname_x1=nickname_x1,
259
+ y_top=y_top,
260
+ y_bottom=y_bottom,
261
+ gutter_px=gutter_px,
262
+ min_roi_width_px=min_roi_width_px,
263
+ )
264
+ else:
265
+ roi_std = _crop_standard_roi_right(
266
+ bgr,
267
+ nickname_x2=nickname_x2,
268
+ y_top=y_top,
269
+ y_bottom=y_bottom,
270
+ gutter_px=gutter_px,
271
+ min_roi_width_px=min_roi_width_px,
272
+ )
273
+ if roi_std is not None and _roi_passes_texture_and_shape(roi_std, **shape_kw):
274
+ return True
275
+ return False
276
+
277
+
278
+ def nickname_row_passes_avatar_roi(
279
+ bgr: Optional[np.ndarray],
280
+ *,
281
+ nickname_x1: int,
282
+ y_top: int,
283
+ y_bottom: int,
284
+ avatar_anchor: Optional["AvatarCentroid"] = None,
285
+ gutter_px: int = 10,
286
+ min_roi_width_px: int = 40,
287
+ laplacian_min: float = 48.0,
288
+ rgb_std_min: float = 4.0,
289
+ edge_sig_cols_min: float = 15.0,
290
+ edge_right_left_ratio_min: float = 1.2,
291
+ prd_strict: bool = False,
292
+ nickname_x2: int = 0,
293
+ side: Side = "left",
294
+ ) -> bool:
295
+ """标准 ROI:纹理∧列缘形状;锚点路径:条带纹理 + 圆盘纹理(几何先验定位头像)。
296
+
297
+ ``prd_strict=True`` 时委托 :func:`nickname_row_passes_prd_avatar_guard`(验收 SSOT)。
298
+ """
299
+ if prd_strict:
300
+ return nickname_row_passes_prd_avatar_guard(
301
+ bgr,
302
+ nickname_x1=nickname_x1,
303
+ nickname_x2=nickname_x2 or nickname_x1,
304
+ y_top=y_top,
305
+ y_bottom=y_bottom,
306
+ side=side,
307
+ avatar_anchor=avatar_anchor,
308
+ gutter_px=gutter_px,
309
+ min_roi_width_px=min_roi_width_px,
310
+ laplacian_min=laplacian_min,
311
+ rgb_std_min=rgb_std_min,
312
+ )
313
+ if cv2 is None or bgr is None or bgr.size == 0:
314
+ return False
315
+
316
+ kw = dict(
317
+ laplacian_min=laplacian_min,
318
+ rgb_std_min=rgb_std_min,
319
+ edge_sig_cols_min=edge_sig_cols_min,
320
+ edge_right_left_ratio_min=edge_right_left_ratio_min,
321
+ )
322
+
323
+ roi_std = _crop_standard_roi(
324
+ bgr,
325
+ nickname_x1=nickname_x1,
326
+ y_top=y_top,
327
+ y_bottom=y_bottom,
328
+ gutter_px=gutter_px,
329
+ min_roi_width_px=min_roi_width_px,
330
+ )
331
+ if roi_std is not None and _roi_passes_texture_and_shape(roi_std, **kw):
332
+ return True
333
+
334
+ if avatar_anchor is None:
335
+ return False
336
+
337
+ roi_strip = _crop_anchor_aligned_roi(
338
+ bgr,
339
+ anchor=avatar_anchor,
340
+ nickname_x1=nickname_x1,
341
+ y_bottom=y_bottom,
342
+ gutter_px=gutter_px,
343
+ min_roi_width_px=min_roi_width_px,
344
+ )
345
+ roi_disk = _crop_avatar_disk_roi(
346
+ bgr,
347
+ anchor=avatar_anchor,
348
+ nickname_x1=nickname_x1,
349
+ gutter_px=gutter_px,
350
+ min_roi_width_px=min_roi_width_px,
351
+ )
352
+ if roi_strip is None or roi_disk is None:
353
+ return False
354
+ tex_kw = dict(laplacian_min=laplacian_min, rgb_std_min=rgb_std_min)
355
+ return _roi_passes_texture(roi_strip, **tex_kw) and _roi_passes_texture(
356
+ roi_disk, **tex_kw
357
+ )
358
+
359
+
360
+ def avatar_roi_pass(
361
+ bgr: Optional[np.ndarray],
362
+ *,
363
+ nickname_x1: int,
364
+ y_top: int,
365
+ y_bottom: int,
366
+ avatar_anchor: Optional["AvatarCentroid"] = None,
367
+ gutter_px: int = 10,
368
+ min_roi_width_px: int = 40,
369
+ laplacian_min: float = 48.0,
370
+ rgb_std_min: float = 4.0,
371
+ edge_sig_cols_min: float = 15.0,
372
+ edge_right_left_ratio_min: float = 1.2,
373
+ require_edge_shape: bool = True,
374
+ strict_narrow_only: bool = False,
375
+ **legacy: Any,
376
+ ) -> bool:
377
+ """``avatar_roi_pass`` 为 ``nickname_row_passes_avatar_roi`` 别名;忽略扩展类旧参数。"""
378
+ _ = legacy
379
+ if strict_narrow_only and avatar_anchor is None:
380
+ pass
381
+ if not require_edge_shape:
382
+ logger.debug("avatar_roi_pass: require_edge_shape=False 已弃用,仍使用纹理∧形状")
383
+ return nickname_row_passes_avatar_roi(
384
+ bgr,
385
+ nickname_x1=nickname_x1,
386
+ y_top=y_top,
387
+ y_bottom=y_bottom,
388
+ avatar_anchor=avatar_anchor,
389
+ gutter_px=gutter_px,
390
+ min_roi_width_px=min_roi_width_px,
391
+ laplacian_min=laplacian_min,
392
+ rgb_std_min=rgb_std_min,
393
+ edge_sig_cols_min=edge_sig_cols_min,
394
+ edge_right_left_ratio_min=edge_right_left_ratio_min,
395
+ )
396
+
397
+
398
+ def nickname_left_roi_passes_avatar_signal(
399
+ bgr: Optional[np.ndarray],
400
+ *,
401
+ nickname_x1: int,
402
+ y_top: int,
403
+ y_bottom: int,
404
+ avatar_anchor: Optional["AvatarCentroid"] = None,
405
+ gutter_px: int = 10,
406
+ min_roi_width_px: int = 40,
407
+ laplacian_min: float = 48.0,
408
+ rgb_std_min: float = 4.0,
409
+ edge_sig_cols_min: float = 15.0,
410
+ edge_right_left_ratio_min: float = 1.2,
411
+ **kwargs: Any,
412
+ ) -> bool:
413
+ _ = kwargs
414
+ return nickname_row_passes_avatar_roi(
415
+ bgr,
416
+ nickname_x1=nickname_x1,
417
+ y_top=y_top,
418
+ y_bottom=y_bottom,
419
+ avatar_anchor=avatar_anchor,
420
+ gutter_px=gutter_px,
421
+ min_roi_width_px=min_roi_width_px,
422
+ laplacian_min=laplacian_min,
423
+ rgb_std_min=rgb_std_min,
424
+ edge_sig_cols_min=edge_sig_cols_min,
425
+ edge_right_left_ratio_min=edge_right_left_ratio_min,
426
+ )
427
+
428
+
429
+ def load_png_bgr(png_path: Any) -> Optional[np.ndarray]:
430
+ if cv2 is None or png_path is None:
431
+ return None
432
+ try:
433
+ return cv2.imread(str(png_path))
434
+ except Exception as e:
435
+ logger.warning("nickname_avatar_guard: imread 失败 path=%s err=%s", png_path, e)
436
+ return None
@@ -0,0 +1,234 @@
1
+ """Optional PaddleOCR pass to read digits off unread badges on the main
2
+ conversation list.
3
+
4
+ Why this module exists
5
+ ----------------------
6
+ The bulk of :func:`TopGroupScanDriver.scan_pinned_groups` is **OCR-free** —
7
+ it only has to tell whether a row has an unread badge (red blob near the
8
+ avatar's top-right corner) and whether the row is pinned (grey cell
9
+ background). The digit count ("3", "12", "99+") is a **nice-to-have**
10
+ used by downstream prioritisation (iterate rows with larger counts
11
+ first) and logging. We keep it behind an opt-in flag so:
12
+
13
+ - The fast path stays fast — PaddleOCR takes ~1.5 s per cold init + a
14
+ few hundred ms per badge, which would slow every live scan loop.
15
+ - Unit tests don't need a working Paddle install to cover the badge
16
+ geometry path.
17
+
18
+ Usage
19
+ -----
20
+ Call :func:`ocr_unread_badge_digits` AFTER ``scan_pinned_groups``
21
+ produced its rows. The function lazily imports PaddleOCR, crops a small
22
+ padded window around each ``UnreadDotHit``, binarises (red-bg/white-fg
23
+ → black-bg/white-fg), upscales 6x, and returns a map
24
+ ``{hit_index → (digits, score)}`` (positional hit index ``0 …``) for every row whose OCR confidently
25
+ matches a digit pattern. Rows without a confident read are absent.
26
+
27
+ Digit regex
28
+ -----------
29
+ ``\\d+\\+?`` — matches "1", "23", "99+", "0" (the last is never shown
30
+ by WeChat but we still accept it to keep the regex portable).
31
+
32
+ Thresholds / pre-processing are the d4-auto-scan calibration values
33
+ that scored ~90%+ on the edb1a89f baseline when the previous scratch
34
+ tool ``tools/detect_unread_badges.py`` (since removed) was still
35
+ around. Kept as module-level constants so Day 7 calibration can bump
36
+ them without touching the driver.
37
+ """
38
+ from __future__ import annotations
39
+
40
+ import logging
41
+ import re
42
+ from typing import Optional
43
+
44
+ import cv2 # type: ignore
45
+ import numpy as np # type: ignore
46
+
47
+ from wechat_screenshot_vision_algorithm.algorithms.template_matching import UnreadDotHit
48
+
49
+ LOG = logging.getLogger("wx_collector.unread_badge_ocr")
50
+
51
+ _DIGIT_RE = re.compile(r"\d+\+?")
52
+
53
+ # Pre-processing ------------------------------------------------------------
54
+
55
+ #: Padding (in pixels AT RAW SCALE) around each badge bbox before cropping.
56
+ #: PaddleOCR's detection head needs ~20% whitespace margin around the glyph.
57
+ #: We take ``max(8, w // 3)`` / ``max(8, h // 3)`` so both 27-px solid dots
58
+ #: and 50-px digit badges get proportionally sized margins.
59
+ PAD_MIN = 8
60
+ PAD_DIV = 3
61
+
62
+ #: Red-mask HSV cut (``cv2.inRange`` low/high). Matches the white-pixel
63
+ #: detector used downstream — pixels that are BRIGHT and DESATURATED get
64
+ #: promoted to white, everything else (red badge body, avatar bleed,
65
+ #: anti-aliased purple) falls to black.
66
+ WHITE_HSV_LO: tuple[int, int, int] = (0, 0, 200)
67
+ WHITE_HSV_HI: tuple[int, int, int] = (179, 60, 255)
68
+
69
+ #: Fallback grey-scale brightness threshold: if HSV mask misses an AA
70
+ #: near-white pixel (V < 200 but > 180 visually white), the classic
71
+ #: ``threshold`` pass recaptures it.
72
+ GRAY_THRESH = 180
73
+
74
+ #: Upscale factor applied to the binarised badge crop. PP-OCRv5 mobile
75
+ #: wants ~32 px min glyph height; a single-digit 27-px badge at 6x
76
+ #: becomes 162 px, well above that floor.
77
+ OCR_UPSCALE = 6.0
78
+
79
+
80
+ def _binarise(crop_bgr: np.ndarray) -> np.ndarray:
81
+ """Convert a red-bg / white-fg badge crop to clean black-bg / white-fg.
82
+
83
+ Two strategies OR'd together:
84
+ 1. HSV ``inRange`` for near-white (low saturation + high value).
85
+ 2. Gray-scale ``threshold`` at ``GRAY_THRESH`` to recapture AA edges.
86
+
87
+ Returns a single-channel ``uint8`` mask. The caller should 3-channel
88
+ expand + cubic-upscale it for PaddleOCR.
89
+ """
90
+ gray = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2GRAY)
91
+ hsv = cv2.cvtColor(crop_bgr, cv2.COLOR_BGR2HSV)
92
+ is_white = cv2.inRange(
93
+ hsv, np.asarray(WHITE_HSV_LO, dtype=np.uint8),
94
+ np.asarray(WHITE_HSV_HI, dtype=np.uint8),
95
+ )
96
+ _, bright = cv2.threshold(gray, GRAY_THRESH, 255, cv2.THRESH_BINARY)
97
+ return cv2.bitwise_or(is_white, bright)
98
+
99
+
100
+ # PaddleOCR init (lazy, cached) --------------------------------------------
101
+
102
+ _OCR_SINGLETON = None
103
+ _OCR_VERSION_MAJOR: Optional[int] = None
104
+
105
+
106
+ def _get_ocr():
107
+ """Lazy PaddleOCR factory. Cached for the process lifetime."""
108
+ global _OCR_SINGLETON, _OCR_VERSION_MAJOR
109
+ if _OCR_SINGLETON is not None:
110
+ return _OCR_SINGLETON, _OCR_VERSION_MAJOR
111
+
112
+ try:
113
+ from paddleocr import PaddleOCR # type: ignore
114
+ import paddleocr as _pkg # type: ignore
115
+ except ImportError as e:
116
+ raise RuntimeError(
117
+ "ocr_unread_badge_digits requires paddleocr; "
118
+ "pip install paddleocr"
119
+ ) from e
120
+
121
+ try:
122
+ ver_major = int(str(getattr(_pkg, "__version__", "2.0.0")).split(".")[0])
123
+ except Exception: # pragma: no cover
124
+ ver_major = 2
125
+
126
+ LOG.info("initialising PaddleOCR v%d (once-per-process)", ver_major)
127
+ if ver_major >= 3:
128
+ ocr = PaddleOCR(
129
+ use_textline_orientation=False,
130
+ lang="ch", device="cpu",
131
+ text_detection_model_name="PP-OCRv5_mobile_det",
132
+ text_recognition_model_name="PP-OCRv5_mobile_rec",
133
+ )
134
+ else: # pragma: no cover — legacy 2.x fallback, project pins 3.x
135
+ ocr = PaddleOCR(use_angle_cls=False, lang="ch", use_gpu=False)
136
+
137
+ _OCR_SINGLETON = ocr
138
+ _OCR_VERSION_MAJOR = ver_major
139
+ return ocr, ver_major
140
+
141
+
142
+ # Public API ---------------------------------------------------------------
143
+
144
+
145
+ def ocr_unread_badge_digits(
146
+ screen_bgr: np.ndarray,
147
+ hits: tuple[UnreadDotHit, ...] | list[UnreadDotHit],
148
+ ) -> dict[int, tuple[str, float]]:
149
+ """Read digit(s) off each badge.
150
+
151
+ Args:
152
+ screen_bgr: The FULL conversation-list screenshot (the same
153
+ ``np.ndarray`` the driver passed to ``detect_unread_dots``).
154
+ We crop from the raw image so the badge stays at native
155
+ resolution before the 6x upscale.
156
+ hits: Sequence of :class:`UnreadDotHit`. Indexing in the
157
+ returned dict is POSITIONAL (``0..len(hits)-1``), NOT
158
+ ``UnreadDotHit`` identity — this keeps the function
159
+ ``dataclass(frozen=True)``-safe.
160
+
161
+ Returns:
162
+ ``{i: (digits, score)}`` for each hit index that produced a
163
+ confident digit match. Hits without a match are absent from
164
+ the dict (rather than mapped to ``None``) so callers can use
165
+ ``dict.get(i)`` freely.
166
+
167
+ Raises:
168
+ RuntimeError: if ``paddleocr`` is not installed.
169
+ """
170
+ if not hits:
171
+ return {}
172
+
173
+ ocr, ver_major = _get_ocr()
174
+ H, W = screen_bgr.shape[:2]
175
+ out: dict[int, tuple[str, float]] = {}
176
+
177
+ for i, b in enumerate(hits):
178
+ pad_x = max(PAD_MIN, b.w // PAD_DIV)
179
+ pad_y = max(PAD_MIN, b.h // PAD_DIV)
180
+ cx1 = max(0, b.x - pad_x)
181
+ cy1 = max(0, b.y - pad_y)
182
+ cx2 = min(W, b.x + b.w + pad_x)
183
+ cy2 = min(H, b.y + b.h + pad_y)
184
+ crop = screen_bgr[cy1:cy2, cx1:cx2]
185
+ if crop.size == 0:
186
+ continue
187
+ binary = _binarise(crop)
188
+ bin_bgr = cv2.merge([binary, binary, binary])
189
+ upscaled = cv2.resize(
190
+ bin_bgr, None, fx=OCR_UPSCALE, fy=OCR_UPSCALE,
191
+ interpolation=cv2.INTER_CUBIC,
192
+ )
193
+ texts: list[str] = []
194
+ scores: list[float] = []
195
+ try:
196
+ if ver_major >= 3:
197
+ results = list(ocr.predict(upscaled))
198
+ if results:
199
+ r = results[0]
200
+ texts = list(r.get("rec_texts", []))
201
+ scores = [float(s) for s in r.get("rec_scores", [])]
202
+ else: # pragma: no cover
203
+ raw = ocr.ocr(upscaled, cls=False)
204
+ if raw and raw[0]:
205
+ texts = [line[1][0] for line in raw[0]]
206
+ scores = [float(line[1][1]) for line in raw[0]]
207
+ except Exception as e: # pragma: no cover — OCR backend failure
208
+ LOG.debug("badge %d OCR failed: %s", i, e)
209
+ continue
210
+
211
+ best_digits: Optional[str] = None
212
+ best_score = 0.0
213
+ for text, score in zip(texts, scores):
214
+ m = _DIGIT_RE.search(text.strip())
215
+ if not m:
216
+ continue
217
+ if score > best_score:
218
+ best_digits = m.group(0)
219
+ best_score = float(score)
220
+ if best_digits is not None:
221
+ out[i] = (best_digits, round(best_score, 3))
222
+ LOG.debug(
223
+ "badge[%d] bbox=(%d,%d,%d,%d) digits=%r score=%.3f",
224
+ i, b.x, b.y, b.w, b.h, best_digits, best_score,
225
+ )
226
+
227
+ LOG.info(
228
+ "ocr_unread_badge_digits: %d/%d badge(s) produced a digit match",
229
+ len(out), len(hits),
230
+ )
231
+ return out
232
+
233
+
234
+ __all__ = ["ocr_unread_badge_digits"]