454 lines
14 KiB
Python
454 lines
14 KiB
Python
import contextlib
|
|
import importlib.util
|
|
import io
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
|
|
SCRIPT_PATH = (
|
|
Path(__file__).resolve().parents[1]
|
|
/ "scripts"
|
|
/ "detect_wcx_title.py"
|
|
)
|
|
SPEC = importlib.util.spec_from_file_location(
|
|
"detect_wcx_title",
|
|
SCRIPT_PATH,
|
|
)
|
|
if SPEC is None or SPEC.loader is None:
|
|
raise RuntimeError(f"Could not load script: {SCRIPT_PATH}")
|
|
|
|
detect_wcx_title = importlib.util.module_from_spec(SPEC)
|
|
SPEC.loader.exec_module(detect_wcx_title)
|
|
|
|
|
|
class PixelClassificationTests(unittest.TestCase):
|
|
def test_nearly_black_thresholds(self) -> None:
|
|
self.assertTrue(
|
|
detect_wcx_title.is_nearly_black(30, 30, 30)
|
|
)
|
|
self.assertFalse(
|
|
detect_wcx_title.is_nearly_black(31, 30, 30)
|
|
)
|
|
|
|
def test_red_thresholds(self) -> None:
|
|
self.assertTrue(
|
|
detect_wcx_title.is_red(150, 100, 100)
|
|
)
|
|
self.assertFalse(
|
|
detect_wcx_title.is_red(139, 0, 0)
|
|
)
|
|
self.assertFalse(
|
|
detect_wcx_title.is_red(150, 101, 100)
|
|
)
|
|
|
|
def test_bright_thresholds(self) -> None:
|
|
self.assertTrue(
|
|
detect_wcx_title.is_bright(180, 180, 180)
|
|
)
|
|
self.assertFalse(
|
|
detect_wcx_title.is_bright(180, 179, 180)
|
|
)
|
|
|
|
def test_layout_metrics_for_constructed_buffer(self) -> None:
|
|
width = 10
|
|
height = 20
|
|
pixels = [(100, 100, 100)] * (width * height)
|
|
|
|
for x_position in range(width):
|
|
pixels[x_position] = (0, 0, 0)
|
|
pixels[width + x_position] = (0, 0, 0)
|
|
|
|
for y_position in range(3):
|
|
pixels[y_position * width] = (150, 0, 0)
|
|
|
|
for y_position in range(height - 3, height):
|
|
pixels[y_position * width + 1] = (255, 255, 255)
|
|
|
|
rgb_data = bytes(
|
|
channel
|
|
for pixel in pixels
|
|
for channel in pixel
|
|
)
|
|
metrics = detect_wcx_title.calculate_layout_metrics(
|
|
rgb_data,
|
|
width=width,
|
|
height=height,
|
|
)
|
|
|
|
self.assertAlmostEqual(
|
|
metrics["black_ratio_top"],
|
|
18 / 20,
|
|
)
|
|
self.assertAlmostEqual(
|
|
metrics["red_ratio_top"],
|
|
3 / 30,
|
|
)
|
|
self.assertAlmostEqual(
|
|
metrics["bright_ratio_bottom"],
|
|
3 / 30,
|
|
)
|
|
|
|
def test_layout_metrics_reject_wrong_buffer_size(self) -> None:
|
|
with self.assertRaisesRegex(ValueError, "Expected 12 RGB bytes"):
|
|
detect_wcx_title.calculate_layout_metrics(
|
|
bytes(11),
|
|
width=2,
|
|
height=2,
|
|
)
|
|
|
|
def test_split_rgb24_frames(self) -> None:
|
|
frame_size = 2 * 2 * 3
|
|
frames = (
|
|
bytes([1]) * frame_size,
|
|
bytes([2]) * frame_size,
|
|
bytes([3]) * frame_size,
|
|
)
|
|
|
|
result = detect_wcx_title.split_rgb24_frames(
|
|
b"".join(frames),
|
|
width=2,
|
|
height=2,
|
|
)
|
|
|
|
self.assertEqual(result, frames)
|
|
|
|
def test_split_rgb24_frames_rejects_short_buffer(self) -> None:
|
|
with self.assertRaisesRegex(ValueError, "too short"):
|
|
detect_wcx_title.split_rgb24_frames(
|
|
bytes(35),
|
|
width=2,
|
|
height=2,
|
|
)
|
|
|
|
def test_split_rgb24_frames_rejects_long_buffer(self) -> None:
|
|
with self.assertRaisesRegex(ValueError, "too long"):
|
|
detect_wcx_title.split_rgb24_frames(
|
|
bytes(37),
|
|
width=2,
|
|
height=2,
|
|
)
|
|
|
|
def test_wcx_layout_score_all_conditions(self) -> None:
|
|
metrics = {
|
|
"black_ratio_total": 0.40,
|
|
"black_ratio_top": 0.60,
|
|
"black_ratio_bottom": 0.70,
|
|
"black_ratio_left": 0.90,
|
|
"black_ratio_right": 0.90,
|
|
"red_ratio_top": 0.10,
|
|
"bright_ratio_top": 0.05,
|
|
"bright_ratio_bottom": 0.08,
|
|
}
|
|
|
|
self.assertEqual(
|
|
detect_wcx_title.calculate_wcx_layout_score(metrics),
|
|
8,
|
|
)
|
|
|
|
def test_wcx_layout_score_no_conditions(self) -> None:
|
|
metrics = {
|
|
"black_ratio_total": 0.10,
|
|
"black_ratio_top": 0.10,
|
|
"black_ratio_bottom": 0.10,
|
|
"black_ratio_left": 0.10,
|
|
"black_ratio_right": 0.10,
|
|
"red_ratio_top": 0.01,
|
|
"bright_ratio_top": 0.01,
|
|
"bright_ratio_bottom": 0.01,
|
|
}
|
|
|
|
self.assertEqual(
|
|
detect_wcx_title.calculate_wcx_layout_score(metrics),
|
|
0,
|
|
)
|
|
|
|
def test_wcx_layout_score_includes_boundaries(self) -> None:
|
|
metrics = {
|
|
"black_ratio_total": 0.20,
|
|
"black_ratio_top": 0.45,
|
|
"black_ratio_bottom": 0.50,
|
|
"black_ratio_left": 0.75,
|
|
"black_ratio_right": 0.75,
|
|
"red_ratio_top": 0.05,
|
|
"bright_ratio_top": 0.02,
|
|
"bright_ratio_bottom": 0.03,
|
|
}
|
|
|
|
self.assertEqual(
|
|
detect_wcx_title.calculate_wcx_layout_score(metrics),
|
|
8,
|
|
)
|
|
|
|
metrics["black_ratio_total"] = 0.60
|
|
self.assertEqual(
|
|
detect_wcx_title.calculate_wcx_layout_score(metrics),
|
|
8,
|
|
)
|
|
|
|
def test_wcs_like_metrics_score_low(self) -> None:
|
|
metrics = {
|
|
"black_ratio_total": 0.88,
|
|
"black_ratio_top": 1.00,
|
|
"black_ratio_bottom": 1.00,
|
|
"black_ratio_left": 0.96,
|
|
"black_ratio_right": 0.96,
|
|
"red_ratio_top": 0.00,
|
|
"bright_ratio_top": 0.00,
|
|
"bright_ratio_bottom": 0.00,
|
|
}
|
|
|
|
self.assertEqual(
|
|
detect_wcx_title.calculate_wcx_layout_score(metrics),
|
|
4,
|
|
)
|
|
|
|
def test_wunf_like_metrics_score_low(self) -> None:
|
|
metrics = {
|
|
"black_ratio_total": 0.15,
|
|
"black_ratio_top": 0.20,
|
|
"black_ratio_bottom": 0.25,
|
|
"black_ratio_left": 0.10,
|
|
"black_ratio_right": 0.10,
|
|
"red_ratio_top": 0.20,
|
|
"bright_ratio_top": 0.01,
|
|
"bright_ratio_bottom": 0.01,
|
|
}
|
|
|
|
self.assertLessEqual(
|
|
detect_wcx_title.calculate_wcx_layout_score(metrics),
|
|
2,
|
|
)
|
|
|
|
def test_select_best_timestamp_uses_highest_score(self) -> None:
|
|
self.assertEqual(
|
|
detect_wcx_title.select_best_timestamp(
|
|
((8.0, 3), (10.0, 8), (12.0, 5))
|
|
),
|
|
(10.0, 8),
|
|
)
|
|
|
|
def test_select_best_timestamp_uses_earliest_tie(self) -> None:
|
|
self.assertEqual(
|
|
detect_wcx_title.select_best_timestamp(
|
|
((12.0, 8), (10.0, 8), (8.0, 7))
|
|
),
|
|
(10.0, 8),
|
|
)
|
|
|
|
def test_classify_wcx_score_eight(self) -> None:
|
|
self.assertEqual(
|
|
detect_wcx_title.classify_wcx_score(8),
|
|
"wcx",
|
|
)
|
|
|
|
def test_classify_wcx_score_seven(self) -> None:
|
|
self.assertEqual(
|
|
detect_wcx_title.classify_wcx_score(7),
|
|
"uncertain",
|
|
)
|
|
|
|
def test_classify_wcx_score_six(self) -> None:
|
|
self.assertEqual(
|
|
detect_wcx_title.classify_wcx_score(6),
|
|
"uncertain",
|
|
)
|
|
|
|
def test_classify_wcx_score_five(self) -> None:
|
|
self.assertEqual(
|
|
detect_wcx_title.classify_wcx_score(5),
|
|
"not_wcx",
|
|
)
|
|
|
|
def test_classify_wcx_score_zero(self) -> None:
|
|
self.assertEqual(
|
|
detect_wcx_title.classify_wcx_score(0),
|
|
"not_wcx",
|
|
)
|
|
|
|
def test_classify_wcx_score_rejects_negative(self) -> None:
|
|
with self.assertRaises(ValueError):
|
|
detect_wcx_title.classify_wcx_score(-1)
|
|
|
|
def test_classify_wcx_score_rejects_above_maximum(self) -> None:
|
|
with self.assertRaises(ValueError):
|
|
detect_wcx_title.classify_wcx_score(9)
|
|
|
|
|
|
class BatchModeTests(unittest.TestCase):
|
|
def test_parse_comma_separated_endings(self) -> None:
|
|
self.assertEqual(
|
|
detect_wcx_title.parse_endings(["mp4,avi"]),
|
|
{"mp4", "avi"},
|
|
)
|
|
|
|
def test_parse_repeated_endings(self) -> None:
|
|
self.assertEqual(
|
|
detect_wcx_title.parse_endings(["mp4", "avi"]),
|
|
{"mp4", "avi"},
|
|
)
|
|
|
|
def test_parse_endings_normalizes_dot_and_case(self) -> None:
|
|
self.assertEqual(
|
|
detect_wcx_title.parse_endings([".MP4"]),
|
|
{"mp4"},
|
|
)
|
|
|
|
def test_video_file_filtering_and_ignored_files(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_directory:
|
|
directory = Path(temporary_directory)
|
|
video = directory / "movie.MP4"
|
|
wrong_ending = directory / "movie.txt"
|
|
apple_double = directory / "._movie.mp4"
|
|
ds_store = directory / ".DS_Store"
|
|
thumbs = directory / "Thumbs.db"
|
|
desktop = directory / "desktop.ini"
|
|
|
|
for path in (
|
|
video,
|
|
wrong_ending,
|
|
apple_double,
|
|
ds_store,
|
|
thumbs,
|
|
desktop,
|
|
):
|
|
path.touch()
|
|
|
|
endings = {"mp4"}
|
|
self.assertTrue(
|
|
detect_wcx_title.is_video_file(video, endings)
|
|
)
|
|
for path in (apple_double, ds_store, thumbs, desktop):
|
|
self.assertTrue(
|
|
detect_wcx_title.is_ignored_file(path)
|
|
)
|
|
for path in (
|
|
wrong_ending,
|
|
apple_double,
|
|
ds_store,
|
|
thumbs,
|
|
desktop,
|
|
):
|
|
self.assertFalse(
|
|
detect_wcx_title.is_video_file(path, endings)
|
|
)
|
|
|
|
def test_non_recursive_file_search(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_directory:
|
|
directory = Path(temporary_directory)
|
|
direct = directory / "direct.mp4"
|
|
nested_directory = directory / "nested"
|
|
nested = nested_directory / "nested.mp4"
|
|
nested_directory.mkdir()
|
|
direct.touch()
|
|
nested.touch()
|
|
|
|
self.assertEqual(
|
|
detect_wcx_title.find_video_files(
|
|
directory,
|
|
endings={"mp4"},
|
|
recursive=False,
|
|
),
|
|
[direct.resolve()],
|
|
)
|
|
|
|
def test_recursive_file_search(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_directory:
|
|
directory = Path(temporary_directory)
|
|
direct = directory / "direct.mp4"
|
|
nested_directory = directory / "nested"
|
|
nested = nested_directory / "nested.mp4"
|
|
nested_directory.mkdir()
|
|
direct.touch()
|
|
nested.touch()
|
|
|
|
self.assertEqual(
|
|
detect_wcx_title.find_video_files(
|
|
directory,
|
|
endings={"mp4"},
|
|
recursive=True,
|
|
),
|
|
sorted(
|
|
[direct.resolve(), nested.resolve()],
|
|
key=lambda path: str(path),
|
|
),
|
|
)
|
|
|
|
def test_file_search_is_sorted_by_full_path(self) -> None:
|
|
with tempfile.TemporaryDirectory() as temporary_directory:
|
|
directory = Path(temporary_directory)
|
|
paths = [
|
|
directory / "z.mp4",
|
|
directory / "A.mp4",
|
|
directory / "m.mp4",
|
|
]
|
|
for path in paths:
|
|
path.touch()
|
|
|
|
result = detect_wcx_title.find_video_files(
|
|
directory,
|
|
endings={"mp4"},
|
|
recursive=False,
|
|
)
|
|
|
|
self.assertEqual(
|
|
result,
|
|
sorted(
|
|
(path.resolve() for path in paths),
|
|
key=lambda path: str(path),
|
|
),
|
|
)
|
|
|
|
def test_batch_counts_classifications(self) -> None:
|
|
video_files = (
|
|
Path("/videos/wcx.mp4"),
|
|
Path("/videos/uncertain.mp4"),
|
|
Path("/videos/not-wcx.mp4"),
|
|
)
|
|
results = {
|
|
video_files[0]: (12.0, 8, "wcx"),
|
|
video_files[1]: (10.0, 7, "uncertain"),
|
|
video_files[2]: (8.0, 4, "not_wcx"),
|
|
}
|
|
|
|
output = io.StringIO()
|
|
with contextlib.redirect_stdout(output):
|
|
exit_code = detect_wcx_title.process_batch(
|
|
video_files,
|
|
lambda path: results[path],
|
|
)
|
|
|
|
self.assertEqual(exit_code, 0)
|
|
self.assertIn("processed: 3", output.getvalue())
|
|
self.assertIn("wcx: 1", output.getvalue())
|
|
self.assertIn("uncertain: 1", output.getvalue())
|
|
self.assertIn("not_wcx: 1", output.getvalue())
|
|
self.assertIn("errors: 0", output.getvalue())
|
|
|
|
def test_batch_returns_one_when_file_errors(self) -> None:
|
|
video_files = (
|
|
Path("/videos/good.mp4"),
|
|
Path("/videos/broken.mp4"),
|
|
)
|
|
|
|
def analyze(path: Path) -> tuple[float, int, str]:
|
|
if path.name == "broken.mp4":
|
|
raise RuntimeError("mock analysis failure")
|
|
return 8.0, 8, "wcx"
|
|
|
|
output = io.StringIO()
|
|
with contextlib.redirect_stdout(output):
|
|
exit_code = detect_wcx_title.process_batch(
|
|
video_files,
|
|
analyze,
|
|
)
|
|
|
|
self.assertEqual(exit_code, 1)
|
|
self.assertIn("error", output.getvalue())
|
|
self.assertIn("processed: 2", output.getvalue())
|
|
self.assertIn("errors: 1", output.getvalue())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|