Initial video AI analysis project
This commit is contained in:
yangyl
167
tests/test_clips.py
Normal file
167
tests/test_clips.py
Normal file
@@ -0,0 +1,167 @@
|
||||
import json
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from video_ai_analysis_poc.clips import build_clip_records, build_clip_records_from_manifest
|
||||
|
||||
|
||||
class ClipTests(unittest.TestCase):
|
||||
def test_build_clip_records_uniformly_samples_frames_per_clip(self):
|
||||
frames = [
|
||||
{
|
||||
"video_id": "video-abc",
|
||||
"frame_id": f"video-abc_f{index + 1:06d}",
|
||||
"frame_path": f"frames/video-abc/{index + 1:06d}.jpg",
|
||||
"offset_seconds": float(index),
|
||||
"timecode": f"00:00:{index:02d}",
|
||||
"pts_time": float(index),
|
||||
"status": "sampled",
|
||||
}
|
||||
for index in range(10)
|
||||
]
|
||||
|
||||
clips = build_clip_records(
|
||||
frames,
|
||||
{
|
||||
"length_seconds": 10,
|
||||
"stride_seconds": 10,
|
||||
"frames_per_clip": 4,
|
||||
"min_frames_per_clip": 2,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(len(clips), 1)
|
||||
self.assertEqual(clips[0]["clip_id"], "video-abc_c000001")
|
||||
self.assertEqual(clips[0]["clip_start_seconds"], 0.0)
|
||||
self.assertEqual(clips[0]["clip_end_seconds"], 10.0)
|
||||
self.assertEqual(
|
||||
[frame["offset_seconds"] for frame in clips[0]["frame_times"]],
|
||||
[0.0, 3.0, 6.0, 9.0],
|
||||
)
|
||||
self.assertEqual(clips[0]["status"], "pending")
|
||||
self.assertEqual(clips[0]["retry_count"], 0)
|
||||
self.assertIsNone(clips[0]["last_error"])
|
||||
|
||||
def test_tail_clip_end_is_truncated_to_last_frame_interval(self):
|
||||
frames = [
|
||||
{
|
||||
"video_id": "video-abc",
|
||||
"frame_id": f"video-abc_f{index + 1:06d}",
|
||||
"frame_path": f"frames/video-abc/{index + 1:06d}.jpg",
|
||||
"offset_seconds": float(index),
|
||||
"timecode": f"00:00:{index:02d}",
|
||||
"pts_time": float(index),
|
||||
"status": "sampled",
|
||||
}
|
||||
for index in range(15)
|
||||
]
|
||||
|
||||
clips = build_clip_records(
|
||||
frames,
|
||||
{
|
||||
"length_seconds": 10,
|
||||
"stride_seconds": 10,
|
||||
"frames_per_clip": 8,
|
||||
"min_frames_per_clip": 4,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(len(clips), 2)
|
||||
self.assertEqual(clips[1]["clip_start_seconds"], 10.0)
|
||||
self.assertEqual(clips[1]["clip_end_seconds"], 15.0)
|
||||
self.assertEqual(clips[1]["clip_end_timecode"], "00:00:15")
|
||||
|
||||
def test_build_clip_records_adds_beijing_time_range_and_frame_times(self):
|
||||
frames = [
|
||||
{
|
||||
"video_id": "video-abc",
|
||||
"frame_id": f"video-abc_f{index + 1:06d}",
|
||||
"frame_path": f"frames/video-abc/{index + 1:06d}.jpg",
|
||||
"offset_seconds": float(index),
|
||||
"timecode": f"00:00:{index:02d}",
|
||||
"pts_time": float(index),
|
||||
"beijing_time": f"2026-06-15 07:00:{index:02d}",
|
||||
"status": "sampled",
|
||||
}
|
||||
for index in range(10)
|
||||
]
|
||||
|
||||
clips = build_clip_records(
|
||||
frames,
|
||||
{
|
||||
"length_seconds": 10,
|
||||
"stride_seconds": 10,
|
||||
"frames_per_clip": 4,
|
||||
"min_frames_per_clip": 2,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(clips[0]["clip_start_beijing_time"], "2026-06-15 07:00:00")
|
||||
self.assertEqual(clips[0]["clip_end_beijing_time"], "2026-06-15 07:00:10")
|
||||
self.assertEqual(
|
||||
[frame["beijing_time"] for frame in clips[0]["frame_times"]],
|
||||
[
|
||||
"2026-06-15 07:00:00",
|
||||
"2026-06-15 07:00:03",
|
||||
"2026-06-15 07:00:06",
|
||||
"2026-06-15 07:00:09",
|
||||
],
|
||||
)
|
||||
|
||||
def test_build_clip_records_from_manifest_skips_failed_frames_and_writes_jsonl(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
root = Path(tmp)
|
||||
frame_manifest = root / "frame_manifest.jsonl"
|
||||
clip_manifest = root / "clip_manifest.jsonl"
|
||||
records = [
|
||||
{
|
||||
"video_id": "video-abc",
|
||||
"frame_id": f"video-abc_f{index + 1:06d}",
|
||||
"frame_path": f"frames/video-abc/{index + 1:06d}.jpg",
|
||||
"offset_seconds": float(index),
|
||||
"timecode": f"00:00:{index:02d}",
|
||||
"pts_time": float(index),
|
||||
"status": "sampled",
|
||||
}
|
||||
for index in range(4)
|
||||
]
|
||||
records.append(
|
||||
{
|
||||
"video_id": "video-abc",
|
||||
"frame_id": None,
|
||||
"frame_path": None,
|
||||
"offset_seconds": None,
|
||||
"timecode": None,
|
||||
"pts_time": None,
|
||||
"status": "sample_failed",
|
||||
"last_error": "bad decode",
|
||||
}
|
||||
)
|
||||
frame_manifest.write_text(
|
||||
"\n".join(json.dumps(record, sort_keys=True) for record in records) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
clips = build_clip_records_from_manifest(
|
||||
frame_manifest,
|
||||
clip_manifest,
|
||||
{
|
||||
"length_seconds": 10,
|
||||
"stride_seconds": 10,
|
||||
"frames_per_clip": 8,
|
||||
"min_frames_per_clip": 4,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(len(clips), 1)
|
||||
self.assertEqual(len(clips[0]["frame_times"]), 4)
|
||||
persisted = [
|
||||
json.loads(line)
|
||||
for line in clip_manifest.read_text(encoding="utf-8").splitlines()
|
||||
]
|
||||
self.assertEqual(persisted, clips)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user