diff --git a/doco/pyproject.toml b/doco/pyproject.toml index 36ffd76..1fafa85 100644 --- a/doco/pyproject.toml +++ b/doco/pyproject.toml @@ -16,6 +16,7 @@ authors = [ dependencies = [ "Pillow>=10.0.0", "imagehash>=4.3.1", + "numpy>=1.24.0", "requests>=2.31.0", "python-dotenv>=1.0.0", "click>=8.1.0", diff --git a/doco/src/doco/video_split.py b/doco/src/doco/video_split.py index a5e6342..7f73b0f 100644 --- a/doco/src/doco/video_split.py +++ b/doco/src/doco/video_split.py @@ -19,6 +19,7 @@ import tempfile from pathlib import Path from typing import Dict, List, Tuple, Optional +import numpy as np from PIL import Image import imagehash @@ -185,14 +186,12 @@ def is_blank_frame(image_path: Path, debug: bool = False) -> Tuple[bool, float]: 返回: (is_blank, white_ratio) """ img = Image.open(image_path).convert("L") - pixels = list(img.getdata()) - total = len(pixels) - white_count = sum(1 for p in pixels if p > BLANK_FRAME_BRIGHTNESS_THRESHOLD) - white_ratio = white_count / total if total > 0 else 0 + arr = np.array(img) + white_ratio = float(np.mean(arr > BLANK_FRAME_BRIGHTNESS_THRESHOLD)) is_blank = white_ratio < BLANK_FRAME_WHITE_PIXEL_RATIO if debug: - frame_idx = image_path.stem # e.g. "frame_0226" + frame_idx = image_path.stem print(f"[debug] {frame_idx}, white_ratio={white_ratio:.6f}, is_blank={is_blank}") return is_blank, white_ratio @@ -226,39 +225,25 @@ def hamming_distance(s1: str, s2: str) -> int: return sum(c1 != c2 for c1, c2 in zip(s1, s2)) -def compute_binary_matrix(image_path: Path, threshold: int = 200) -> List[List[int]]: +def compute_binary_matrix(image_path: Path, threshold: int = 200) -> np.ndarray: """ 将图片转为二值化矩阵(亮度 > threshold → 1,否则 → 0) 用于 IoU 对比 + + 返回: np.ndarray (dtype=uint8, 值为 0 或 1) """ img = Image.open(image_path).convert("L") - w, h = img.size - matrix = [] - for y in range(h): - row = [] - for x in range(w): - pixel = img.getpixel((x, y)) - row.append(1 if pixel > threshold else 0) - matrix.append(row) - return matrix + arr = np.array(img) + return (arr > threshold).astype(np.uint8) -def compute_iou(matrix1: List[List[int]], matrix2: List[List[int]]) -> float: +def compute_iou(matrix1: np.ndarray, matrix2: np.ndarray) -> float: """ 计算两个二值化矩阵的 IoU(交并比) + matrix1, matrix2: np.ndarray (dtype=uint8, 值为 0 或 1) """ - h = len(matrix1) - w = len(matrix1[0]) if h > 0 else 0 - intersection = 0 - union = 0 - for y in range(h): - for x in range(w): - v1 = matrix1[y][x] - v2 = matrix2[y][x] - if v1 == 1 or v2 == 1: - union += 1 - if v1 == 1 and v2 == 1: - intersection += 1 + intersection = int(np.sum((matrix1 == 1) & (matrix2 == 1))) + union = int(np.sum((matrix1 == 1) | (matrix2 == 1))) return intersection / union if union > 0 else 1.0