yingshiyun/analysis/timestamp_reader.py

"""从摄像头画面自带的OSD时间戳水印读取精确到秒的录制时间。

原理：画面固定位置烧录 "YYYY-MM-DD HH:MM:SS" 字样。由于点阵字体渲染时字符间
不保证有空白列（会粘连），不用连通域/空白列分割单个字符，而是先定位"日期组"
"时间组"两个文本块的整体边界（这两组之间的空格间隙远大于组内字符间隙），再用
已知的固定字符数（日期10个字符，时间8个字符）做等宽硬切割。

日期组、时间组分别独立做 Otsu 二值化（而非整行统一阈值）：两组背景光照条件
可能不同，统一阈值会让偏暗的一组文字被过度清除。

数字识别用多模板 IOU（交并比）相似度匹配，而非 cv2.matchTemplate 的
TM_CCOEFF_NORMED —— 二值图像中大量近空白/全黑的 cell 两两比较时，归一化相关
系数会退化出虚假的高分。模板库按分辨率持久化在 analysis/templates/<分辨率>/
digit_templates.npz，收集自多个已知真值样本中同一数字在所有出现位置的实例。
"""
from datetime import datetime, timedelta

import cv2
import numpy as np

from analysis.config import PROJECT_ROOT

TEMPLATES_DIR = PROJECT_ROOT / "analysis" / "templates"

PATTERN = ("digit digit digit digit dash digit digit dash digit digit "
           "digit digit colon digit digit colon digit digit").split()

# 每种分辨率的裁剪坐标 + 日期组/时间组边界（相对裁剪区域的局部坐标）+ cell 尺寸。
# 边界均由列投影法（找字符间隙）在多个样本上标定校正，不是拍脑袋的估计值。
RESOLUTION_PROFILES = {
    (2560, 1440): {
        "crop_y": (1372, 1402),
        "crop_x": (1750, 2560),
        "date_group": (108, 373),   # 10字符: Y Y Y Y - M M - D D
        "time_group": (407, 613),   # 8字符: H H : M M : S S
        "cell_size": (20, 30),      # (width, height)，传给 cv2.resize
    },
    (1280, 720): {
        "crop_y": (682, 705),
        "crop_x": (870, 1280),
        "date_group": (60, 192),
        "time_group": (208, 312),
        "cell_size": (20, 30),
    },
}

# 识别置信度低于此值时判定为不可信，调用方应回退到"只精确到天"的粗粒度时间
MIN_CONFIDENCE = 0.5

_template_cache = {}


def _load_templates(profile_key):
    """按分辨率加载模板库，进程内缓存一次。"""
    if profile_key not in _template_cache:
        npz_path = TEMPLATES_DIR / f"{profile_key[1]}p" / "digit_templates.npz"
        data = np.load(npz_path)
        _template_cache[profile_key] = {digit: data[digit] for digit in data.files}
    return _template_cache[profile_key]


def _cells_from_group(gray, group, n, cell_size):
    s, e = group
    sub = gray[:, s:e]
    _, thresh = cv2.threshold(sub, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    w = (e - s) / n
    cells = []
    for i in range(n):
        x0 = int(round(i * w))
        x1 = int(round((i + 1) * w))
        cells.append(cv2.resize(thresh[:, x0:x1], cell_size, interpolation=cv2.INTER_AREA))
    return cells


def _iou(a, b):
    a = a > 0
    b = b > 0
    inter = np.logical_and(a, b).sum()
    union = np.logical_or(a, b).sum()
    return inter / union if union > 0 else 0.0


def _match_digit(cell, templates):
    best_ch, best_score = None, -1
    for digit, tmpls in templates.items():
        for tmpl in tmpls:
            score = _iou(cell, tmpl)
            if score > best_score:
                best_score, best_ch = score, digit
    return best_ch, best_score


def _decode(gray, profile, templates):
    cells = (_cells_from_group(gray, profile["date_group"], 10, profile["cell_size"])
             + _cells_from_group(gray, profile["time_group"], 8, profile["cell_size"]))

    out = []
    scores = []
    for cell, kind in zip(cells, PATTERN):
        if kind == "digit":
            ch, score = _match_digit(cell, templates)
            out.append(ch)
            scores.append(score)
        elif kind == "dash":
            out.append("-")
        elif kind == "colon":
            out.append(":")

    s = "".join(out)
    # out 顺序 = 日期组10个cell(已含2个'-') + 时间组8个cell(已含2个':')
    formatted = f"{s[0:10]} {s[10:18]}"
    return formatted, (min(scores) if scores else 0.0)


def read_timestamp(file_path, capture_date=None):
    """读取视频首帧的OSD时间戳水印，返回 "YYYY-MM-DD HH:MM:SS" 字符串。

    分辨率不在已标定范围内、读帧失败、识别置信度过低、或识别结果未通过合法性
    校验时返回 None，调用方应保留原有"只精确到天"的粗粒度 start_time 作为兜底。

    单纯依赖 IOU 置信度阈值并不可靠：抽样发现即使分数达标，仍会出现"分钟81"
    这类不合法数值、或年份被识别错一位数（如2026错认成2006）但恰好凑成一个
    合法日期的情况。因此额外做两层校验：
    1. 用 datetime 解析，过滤掉月/日/时/分/秒超出合法范围的结果
    2. 若调用方提供 capture_date（文件所在日期文件夹，格式YYYYMMDD），解码出的
       日期必须与之相差在1天以内（凌晨录制可能跨天），否则判定为误识别
    """
    cap = cv2.VideoCapture(str(file_path))
    ok, frame = cap.read()
    cap.release()
    if not ok:
        return None

    height, width = frame.shape[:2]
    profile = RESOLUTION_PROFILES.get((width, height))
    if profile is None:
        return None

    y0, y1 = profile["crop_y"]
    x0, x1 = profile["crop_x"]
    gray = cv2.cvtColor(frame[y0:y1, x0:x1], cv2.COLOR_BGR2GRAY)

    templates = _load_templates((width, height))
    formatted, min_score = _decode(gray, profile, templates)
    if min_score < MIN_CONFIDENCE:
        return None

    try:
        parsed = datetime.strptime(formatted, "%Y-%m-%d %H:%M:%S")
    except ValueError:
        return None

    if capture_date is not None:
        folder_date = datetime.strptime(capture_date, "%Y%m%d")
        if abs((parsed.date() - folder_date.date())) > timedelta(days=1):
            return None

    return formatted