File size: 5,320 Bytes
e0db531 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | """Stage 2 of the VIPER-19K pipeline: three-axis physics annotation.
The paper annotates every surviving clip on three axes -- **material**,
**trajectory** and **physical impact**. Those labels are what the pairing stage
buckets on, so they are drawn from a closed taxonomy (free-form text would not
bucket). We additionally keep a one-sentence physics summary, which is used as
the target prompt during training.
Run with ``--shard i --num_shards N`` to spread over GPUs.
"""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
import numpy as np
import torch
MATERIALS = [
"rigid_solid", "elastic_deformable", "plastic_deformable", "brittle",
"granular", "liquid", "viscous_liquid", "soft_body_cloth", "gas_smoke",
]
TRAJECTORIES = [
"free_fall", "projectile_arc", "rolling", "sliding", "swinging_pendulum",
"rotating_spinning", "flowing_pouring", "oscillating_vibrating",
"expanding_dispersing", "in_place_deformation",
]
IMPACTS = [
"collision_impact", "compression_squeeze", "fracture_shatter",
"dent_deformation", "bounce_rebound", "splash_dispersal",
"melting_phase_change", "tearing_rupture", "crumbling", "no_impact",
]
PROMPT = f"""Analyse the physical process in this video and answer STRICTLY as JSON.
Choose exactly one value for each axis from these closed sets:
material: {MATERIALS}
trajectory: {TRAJECTORIES}
physical_impact: {IMPACTS}
Return this JSON object and nothing else:
{{
"material": "<one of material>",
"trajectory": "<one of trajectory>",
"physical_impact": "<one of physical_impact>",
"physics_summary": "<one sentence describing the physical dynamics: how the \
object moves, deforms and interacts. Describe the PHYSICS, not the appearance \
or identity of the objects.>"
}}"""
def sample_frames(path: str, n: int = 8) -> np.ndarray:
import imageio.v3 as iio
v = iio.imread(path, plugin="pyav") # [T,H,W,3] uint8
idx = np.linspace(0, len(v) - 1, n).astype(int)
return v[idx]
def parse_json(text: str) -> dict | None:
text = text.strip()
if "```" in text:
text = text.split("```")[1]
if text.startswith("json"):
text = text[4:]
s, e = text.find("{"), text.rfind("}")
if s < 0 or e < 0:
return None
try:
return json.loads(text[s:e + 1])
except json.JSONDecodeError:
return None
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--clips", default="data/filtered_clips.jsonl")
ap.add_argument("--out", default="data/annotated.jsonl")
ap.add_argument("--mllm", default="models/Qwen3-VL-4B-Instruct")
ap.add_argument("--num_frames", type=int, default=8)
ap.add_argument("--shard", type=int, default=0)
ap.add_argument("--num_shards", type=int, default=1)
ap.add_argument("--limit", type=int, default=0)
args = ap.parse_args()
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration
rows = [json.loads(l) for l in open(args.clips)]
if args.limit:
rows = rows[: args.limit]
rows = rows[args.shard::args.num_shards]
proc = AutoProcessor.from_pretrained(args.mllm)
model = Qwen3VLForConditionalGeneration.from_pretrained(
args.mllm, dtype=torch.bfloat16, device_map="cuda",
attn_implementation="sdpa",
).eval()
out_path = args.out if args.num_shards == 1 else f"{args.out}.{args.shard}"
Path(out_path).parent.mkdir(parents=True, exist_ok=True)
n_ok = 0
with open(out_path, "w") as f:
for i, r in enumerate(rows):
try:
frames = sample_frames(r["video"], args.num_frames)
except Exception as e:
print(f"[skip decode] {r['id']}: {e}", flush=True)
continue
messages = [{"role": "user", "content": [
{"type": "video"}, {"type": "text", "text": PROMPT}]}]
text = proc.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True)
inputs = proc(text=[text], videos=[frames], return_tensors="pt").to("cuda")
with torch.inference_mode():
gen = model.generate(**inputs, max_new_tokens=256, do_sample=False)
reply = proc.batch_decode(
gen[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)[0]
ann = parse_json(reply)
if not ann or not all(k in ann for k in
("material", "trajectory", "physical_impact")):
print(f"[skip parse] {r['id']}", flush=True)
continue
# Reject hallucinated out-of-taxonomy values.
if (ann["material"] not in MATERIALS
or ann["trajectory"] not in TRAJECTORIES
or ann["physical_impact"] not in IMPACTS):
print(f"[skip taxonomy] {r['id']}: {ann}", flush=True)
continue
r.update(ann)
f.write(json.dumps(r, ensure_ascii=False) + "\n")
f.flush()
n_ok += 1
if i % 25 == 0:
print(f"[{args.shard}] {i}/{len(rows)} ok={n_ok}", flush=True)
print(f"[{args.shard}] DONE ok={n_ok}/{len(rows)} -> {out_path}", flush=True)
if __name__ == "__main__":
main()
|