simonkoson
206 lines
6.4 KiB
Python
206 lines
6.4 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[:100, :10] = 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 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):
|
|
"""white_ratio>=0.005 但 max_brightness<240 → 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阈值),其余黑
|
|
arr[:100, :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.005
|