simonkoson
207 lines
6.5 KiB
Python
207 lines
6.5 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
video_split 单元测试
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
from pathlib import Path
|
||
|
||
import pytest
|
||
|
||
# 确保 src 在 path 中
|
||
import sys
|
||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||
|
||
from doco.video_split import (
|
||
hamming_distance,
|
||
format_timestamp,
|
||
build_b_manuscript,
|
||
is_blank_frame,
|
||
)
|
||
|
||
|
||
class TestHammingDistance:
|
||
def test_identical_strings(self):
|
||
assert hamming_distance("abc", "abc") == 0
|
||
|
||
def test_different_strings(self):
|
||
assert hamming_distance("abc", "abd") == 1
|
||
|
||
def test_different_length(self):
|
||
# 长度不同时,返回较长字符串的长度
|
||
assert hamming_distance("abc", "ab") == 3
|
||
|
||
|
||
class TestFormatTimestamp:
|
||
def test_zero(self):
|
||
assert format_timestamp(0) == "0m0s"
|
||
|
||
def test_seconds_only(self):
|
||
assert format_timestamp(30000) == "0m30s" # 30秒
|
||
|
||
def test_minutes_and_seconds(self):
|
||
assert format_timestamp(90000) == "1m30s" # 1分30秒
|
||
|
||
def test_longer(self):
|
||
assert format_timestamp(3723000) == "62m3s" # 62分3秒
|
||
|
||
|
||
class TestBuildBManuscript:
|
||
def test_empty_keyframes(self):
|
||
lines = build_b_manuscript([])
|
||
assert lines == []
|
||
|
||
def test_single_frame(self):
|
||
keyframes = [
|
||
{"timestamp_ms": 0, "ocr_text": "测试字幕"}
|
||
]
|
||
lines = build_b_manuscript(keyframes)
|
||
assert len(lines) == 1
|
||
assert "[0m0s]" in lines[0]
|
||
assert "测试字幕" in lines[0]
|
||
|
||
def test_duplicate_text_merged(self):
|
||
"""相邻同文本应合并"""
|
||
keyframes = [
|
||
{"timestamp_ms": 0, "ocr_text": "相同"},
|
||
{"timestamp_ms": 1000, "ocr_text": "相同"},
|
||
{"timestamp_ms": 2000, "ocr_text": "不同"},
|
||
]
|
||
lines = build_b_manuscript(keyframes)
|
||
assert len(lines) == 2
|
||
|
||
def test_placeholder_skipped(self):
|
||
"""OCR占位文本应跳过"""
|
||
keyframes = [
|
||
{"timestamp_ms": 0, "ocr_text": "[OCR待填充 frame=001.png]"},
|
||
{"timestamp_ms": 1000, "ocr_text": "真实字幕"},
|
||
]
|
||
lines = build_b_manuscript(keyframes)
|
||
assert len(lines) == 1
|
||
assert "真实字幕" in lines[0]
|
||
|
||
|
||
class TestKeyframesJson:
|
||
"""验证 keyframes.json 输出格式"""
|
||
|
||
def test_keyframe_structure(self, tmp_path):
|
||
"""验证单个关键帧的 JSON 结构"""
|
||
# 模拟关键帧数据
|
||
kf = {
|
||
"frame_index": 1,
|
||
"timestamp_ms": 1000,
|
||
"frame_image_path": str(tmp_path / "frame_0001.png"),
|
||
"phash": "ff00aabb12345678",
|
||
"ocr_text": "测试",
|
||
}
|
||
|
||
# 验证字段存在
|
||
assert "frame_index" in kf
|
||
assert "timestamp_ms" in kf
|
||
assert "frame_image_path" in kf
|
||
assert "phash" in kf
|
||
assert "ocr_text" in kf
|
||
|
||
def test_keyframes_json_output(self, tmp_path):
|
||
"""验证完整 keyframes.json 输出"""
|
||
frames_dir = tmp_path / "frames"
|
||
frames_dir.mkdir()
|
||
|
||
# 创建一个测试图片
|
||
test_img = frames_dir / "frame_0001.png"
|
||
test_img.write_bytes(b"fake_png_data")
|
||
|
||
keyframes_data = {
|
||
"video_path": str(tmp_path / "video.mp4"),
|
||
"fps_sampled": 1,
|
||
"phash_threshold": 8,
|
||
"keyframes": [
|
||
{
|
||
"frame_index": 0,
|
||
"timestamp_ms": 0,
|
||
"frame_image_path": str(test_img),
|
||
"phash": "abcd1234",
|
||
"ocr_text": "首帧",
|
||
}
|
||
]
|
||
}
|
||
|
||
json_path = tmp_path / "keyframes.json"
|
||
with open(json_path, "w", encoding="utf-8") as f:
|
||
json.dump(keyframes_data, f, ensure_ascii=False, indent=2)
|
||
|
||
# 验证可读
|
||
with open(json_path, "r", encoding="utf-8") as f:
|
||
loaded = json.load(f)
|
||
|
||
assert loaded["fps_sampled"] == 1
|
||
assert loaded["phash_threshold"] == 8
|
||
assert len(loaded["keyframes"]) == 1
|
||
assert loaded["keyframes"][0]["frame_index"] == 0
|
||
|
||
|
||
class TestIsBlankFrame:
|
||
"""is_blank_frame 双条件检测单元测试"""
|
||
|
||
def test_pure_black_frame(self, tmp_path):
|
||
"""全 0 纯黑图:max_brightness=0,white_ratio=0,is_blank=True"""
|
||
from PIL import Image
|
||
img_path = tmp_path / "black.png"
|
||
img = Image.new("L", (100, 100), 0) # 全黑
|
||
img.save(img_path)
|
||
|
||
is_blank, white_ratio, max_brightness = is_blank_frame(img_path)
|
||
assert is_blank is True
|
||
assert max_brightness == 0
|
||
assert white_ratio == 0.0
|
||
|
||
def test_subtitle_frame(self, tmp_path):
|
||
"""有少量白像素(亮度255,占~1%)的图:is_blank=False"""
|
||
from PIL import Image
|
||
import numpy as np
|
||
img_path = tmp_path / "subtitle.png"
|
||
arr = np.zeros((100, 100), dtype=np.uint8)
|
||
# 1% 白像素(亮度255),其余黑(亮度0)
|
||
arr[:10, :10] = 255 # 10×10=100/10000=1%
|
||
img = Image.fromarray(arr, mode="L")
|
||
img.save(img_path)
|
||
|
||
is_blank, white_ratio, max_brightness = is_blank_frame(img_path)
|
||
assert is_blank is False
|
||
assert max_brightness == 255
|
||
assert 0.005 <= white_ratio <= 0.02 # ~1%
|
||
|
||
def test_max_brightness_ok_but_ratio_too_low(self, tmp_path):
|
||
"""max_brightness>=240 但 white_ratio<0.005 → is_blank=True"""
|
||
from PIL import Image
|
||
import numpy as np
|
||
img_path = tmp_path / "few_pixels.png"
|
||
arr = np.zeros((100, 100), dtype=np.uint8)
|
||
# 只有 0.1% 白像素(不够 0.5%),但 max_brightness=255
|
||
arr[:10, :1] = 255
|
||
img = Image.fromarray(arr, mode="L")
|
||
img.save(img_path)
|
||
|
||
is_blank, white_ratio, max_brightness = is_blank_frame(img_path)
|
||
assert is_blank is True
|
||
assert max_brightness == 255
|
||
assert white_ratio < 0.005
|
||
|
||
def test_ratio_ok_but_max_brightness_too_low(self, tmp_path):
|
||
"""max_brightness<240 → 白像素计数全为0 → is_blank=True"""
|
||
from PIL import Image
|
||
import numpy as np
|
||
img_path = tmp_path / "dim_pixels.png"
|
||
arr = np.zeros((100, 100), dtype=np.uint8)
|
||
# 1% 像素亮度=220(不满足阈值240),其余黑
|
||
# is_blank_frame 用 arr > 240 统计白像素,220 全不满足 → white_ratio=0
|
||
arr[:10, :10] = 220
|
||
img = Image.fromarray(arr, mode="L")
|
||
img.save(img_path)
|
||
|
||
is_blank, white_ratio, max_brightness = is_blank_frame(img_path)
|
||
assert is_blank is True
|
||
assert max_brightness == 220
|
||
assert white_ratio == 0.0
|