| |
| """Unit and CPU smoke checks for fixed-frame SegmentQA preprocessing.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| import shutil |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| import av |
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(PROJECT_ROOT / "scripts")) |
|
|
| import preprocess_segmentqa_fixed_frames as prep |
|
|
|
|
| PYTHON = "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/conda-envs/medrl/bin/python" |
| DATA_ROOT = "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge" |
| OUT = PROJECT_ROOT / "tmp/segmentqa_fixed80_preprocess_smoke_20260826" |
|
|
|
|
| def test_env() -> dict[str, str]: |
| env = dict(os.environ) |
| env["PYTHONPATH"] = str(PROJECT_ROOT / "ms-swift") |
| return env |
|
|
|
|
| def video_stats(path: Path) -> dict: |
| with av.open(str(path)) as container: |
| stream = container.streams.video[0] |
| pts = [] |
| dims = set() |
| for frame in container.decode(stream): |
| pts.append(float(frame.pts * frame.time_base)) |
| dims.add((frame.width, frame.height)) |
| return {"count": len(pts), "pts": pts, "dims": dims} |
|
|
|
|
| def test_time_and_sampling_helpers() -> None: |
| assert prep.parse_time("01:02:03") == 3723 |
| assert prep.parse_time(12.5) == 12.5 |
| targets = prep.target_times(10, 18, 4) |
| assert targets == [11, 13, 15, 17] |
| assert prep.metadata_from_pts([10.1234, 10.5678], 420, 224) == { |
| "fps": 1000, |
| "frames_indices": [10123, 10568], |
| "total_num_frames": 10569, |
| "width": 420, |
| "height": 224, |
| "duration": 10.569, |
| "video_backend": "precomputed_absolute_pts", |
| } |
|
|
|
|
| def test_template_metadata_override_path() -> None: |
| from swift.template.templates.qwen import _extract_precomputed_absolute_video_metadata as extract_metadata |
|
|
| metadata = prep.metadata_from_pts([1.0, 2.0], 420, 224) |
| video_inputs = {"video": "clip.mp4", "nframes": 80, "enable_thinking": False, "video_metadata": metadata} |
| extracted = extract_metadata(video_inputs, 0) |
| assert extracted == metadata |
| assert video_inputs == {"video": "clip.mp4", "nframes": 80, "enable_thinking": False} |
| legacy_inputs = {"video": "clip.mp4", "fps": 1.0, "min_frames": 16, "max_frames": 64} |
| assert extract_metadata(legacy_inputs, 0) is None |
| assert legacy_inputs == {"video": "clip.mp4", "fps": 1.0, "min_frames": 16, "max_frames": 64} |
|
|
|
|
| def validate_qwen_fetch_video_path(record: dict) -> None: |
| from qwen_vl_utils import fetch_video |
| from swift.template.templates.qwen import _extract_precomputed_absolute_video_metadata |
| from transformers.video_utils import VideoMetadata |
|
|
| video_inputs = {"video": record["videos"][0], **record["chat_template_kwargs"]} |
| metadata = _extract_precomputed_absolute_video_metadata(video_inputs, 0) |
| assert metadata is not None |
| assert "do_sample_frames" not in metadata |
| VideoMetadata(**metadata) |
| (video, fetched_metadata), sample_fps = fetch_video( |
| video_inputs, return_video_sample_fps=True, return_video_metadata=True, image_patch_size=16) |
| assert video.shape[0] == 80 |
| assert video.shape[1] == 3 |
| assert video.shape[-2] % 32 == 0 |
| assert video.shape[-1] % 32 == 0 |
| assert fetched_metadata["total_num_frames"] >= 80 |
| assert sample_fps > 0 |
| assert video_inputs == {"video": record["videos"][0], "nframes": 80, "enable_thinking": False} |
|
|
|
|
| def assert_author_direct_segment_prompt(prompt: str) -> None: |
| assert "timestamp in the source video" in prompt |
| assert "<1234.5 seconds> is 00:20:34" in prompt |
| assert "START OF THE VIDEO, not from the start of this" in prompt |
| assert ( |
| "The foreign object classes are exactly: Sponge, Clip, Specimen Bag, Silicone Loop, External Drain, " |
| "Needle, Gallstone, Specimen, Mesh, Absorbable Hemostatic Agent." |
| ) in prompt |
| assert "Reply with the answer and nothing else -- no reasoning, no preamble, no" in prompt |
| assert "A single short line." in prompt |
|
|
|
|
| def test_cpu_smoke_preprocess_train() -> None: |
| if OUT.exists(): |
| shutil.rmtree(OUT) |
| cmd = [ |
| PYTHON, |
| "scripts/preprocess_segmentqa_fixed_frames.py", |
| "--data-root", |
| DATA_ROOT, |
| "--output-dir", |
| str(OUT), |
| "--splits", |
| "train", |
| "--frames", |
| "80", |
| "--width", |
| "420", |
| "--height", |
| "224", |
| "--workers", |
| "1", |
| "--limit", |
| "2", |
| ] |
| result = subprocess.run(cmd, cwd=PROJECT_ROOT, check=True, text=True, capture_output=True, env=test_env()) |
| summary = json.loads(result.stdout.strip().splitlines()[-1]) |
| assert summary["segments"] == 2 |
| assert summary["success"] == 2 |
| assert summary["failed"] == 0 |
|
|
| manifest = [json.loads(line) for line in (OUT / "segment_manifest.jsonl").read_text().splitlines()] |
| assert {item["dataset"] for item in manifest} == set(prep.DATASET_NAMES) |
| for item in manifest: |
| stats = video_stats(OUT / item["relpath"]) |
| assert stats["count"] == 80 |
| assert stats["dims"] == {(420, 224)} |
| assert item["verified_frame_count"] == 80 |
| assert item["width"] == 420 |
| assert item["height"] == 224 |
| assert len(item["selected_source_pts"]) == 80 |
| assert all(item["start_seconds"] <= pts < item["end_seconds"] for pts in item["selected_source_pts"]) |
| assert item["video_metadata"]["fps"] == 1000 |
| assert item["video_metadata"]["frames_indices"] == [round(pts * 1000) for pts in item["selected_source_pts"]] |
| assert item["video_metadata"]["total_num_frames"] == max(item["video_metadata"]["frames_indices"]) + 1 |
| assert item["video_metadata"]["video_backend"] == "precomputed_absolute_pts" |
| assert "do_sample_frames" not in item["video_metadata"] |
|
|
| train_rows = [json.loads(line) for line in (OUT / "train.jsonl").read_text().splitlines()] |
| assert len(train_rows) >= 2 |
| assert all(row["messages"][0]["role"] == "system" for row in train_rows) |
| assert_author_direct_segment_prompt(train_rows[0]["messages"][0]["content"]) |
| assert all(row["messages"][1]["content"].startswith("<video>\n") for row in train_rows) |
| assert all(row["messages"][2]["role"] == "assistant" and row["messages"][2]["content"] for row in train_rows) |
| assert all(row["chat_template_kwargs"]["nframes"] == 80 for row in train_rows) |
| assert all(row["chat_template_kwargs"]["enable_thinking"] is False for row in train_rows) |
| validate_qwen_fetch_video_path(train_rows[0]) |
| assert not (OUT / "test_requests.jsonl").exists() |
| report = json.loads((OUT / "validation_report.json").read_text()) |
| assert report["status"] == "pass" |
| assert report["success"] == report["segments"] == 2 |
|
|
| before = {p: p.stat().st_mtime for p in OUT.glob("videos/train/*/*.mp4")} |
| resume = subprocess.run(cmd + ["--resume"], cwd=PROJECT_ROOT, check=True, text=True, capture_output=True, env=test_env()) |
| resume_summary = json.loads(resume.stdout.strip().splitlines()[-1]) |
| assert resume_summary["success"] == 2 |
| assert before == {p: p.stat().st_mtime for p in OUT.glob("videos/train/*/*.mp4")} |
|
|
|
|
| if __name__ == "__main__": |
| test_time_and_sampling_helpers() |
| test_template_metadata_override_path() |
| test_cpu_smoke_preprocess_train() |
|
|
