| """ |
| Caption humanoid singleview training videos (makovian / non_makovian) with |
| Qwen3-VL-30B-A3B-Instruct-FP8 and write DreamGen-ready flat structure. |
| |
| For each video in datasets/humanoid/singleview/{makovian,non_makovian}/videos/...: |
| 1. Resolve symlink → identify source (GR1_robot or DreamDojo-HV_Eval). |
| 2. Fetch the existing task-text hint from gr1robot_task_labels.csv (GR1_robot) |
| or gr1_task_labels_cache.jsonl (DreamDojo). |
| 3. Sample N frames from the video + prepend the task hint in the prompt. |
| 4. Run Qwen3-VL-30B to produce a natural-language caption. |
| |
| Outputs (written in parallel): |
| a) datasets/humanoid/singleview/{split}/captions/<stem>.txt |
| → one-line Qwen caption per episode (intermediate cache). |
| b) datasets/humanoid_dreamgen/ |
| videos/<split>__<stem>.mp4 (symlink) |
| metas/<split>__<stem>.txt (Qwen caption) |
| → flat DreamGen training structure, ready for compute_t5_embeddings.sh. |
| |
| Usage (from project root): |
| python scripts/caption_humanoid_singleview.py |
| python scripts/caption_humanoid_singleview.py --splits makovian --dry_run |
| python scripts/caption_humanoid_singleview.py --overwrite --num_video_frames 12 |
| |
| Prerequisites: |
| pip install vllm>=0.8 transformers>=4.57 imageio imageio-ffmpeg pillow |
| (run from a venv with the above, NOT the DreamDojo venv) |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import json |
| import os |
| import re |
| from collections import defaultdict |
| from pathlib import Path |
|
|
| import imageio.v3 as iio |
| from PIL import Image |
|
|
| os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" |
|
|
| |
| |
| |
| CAMERA = "observation.images.ego_view_freq20" |
|
|
| CAPTION_PROMPT = ( |
| "You are labelling a short humanoid robot tele-operation clip for a " |
| "video-generation model.\n" |
| "The robot's task is: \"{task_hint}\"\n" |
| "Looking at the provided video frames, write ONE concise English sentence " |
| "(max 30 words) describing what the humanoid robot is doing — mention which " |
| "arm/hand is used, the manipulated object(s), and the action verb. " |
| "Do NOT describe background, camera angle, or lighting. " |
| "Output the sentence only, no quotes, no prefix." |
| ) |
|
|
| CAPTION_PROMPT_NO_HINT = ( |
| "You are labelling a short humanoid robot tele-operation clip for a " |
| "video-generation model. In ONE concise English sentence (max 30 words) " |
| "describe what the humanoid robot is doing — which arm/hand is used, the " |
| "manipulated object(s), and the action verb. Do NOT describe background, " |
| "camera angle, or lighting. Output the sentence only, no quotes, no prefix." |
| ) |
|
|
| WS = Path(__file__).resolve().parents[1] |
|
|
|
|
| |
| |
| |
|
|
| def load_gr1robot_task_hints(ws: Path) -> dict[int, str]: |
| """episode_index -> clean task text (no 'locked waist: ' prefix).""" |
| ep_file = ws / "datasets/humanoid/singleview/PhysicalAI-Robotics-GR00T-Teleop-GR1/GR1_robot/meta/episodes.jsonl" |
| ep_to_task: dict[int, str] = {} |
| if not ep_file.exists(): |
| return ep_to_task |
| with ep_file.open() as f: |
| for line in f: |
| line = line.strip() |
| if not line: |
| continue |
| obj = json.loads(line) |
| raw = (obj.get("tasks") or [""])[0] |
| clean = raw.replace("locked waist: ", "").strip() |
| ep_to_task[obj["episode_index"]] = clean |
| return ep_to_task |
|
|
|
|
| def load_dreamdojo_task_hints(ws: Path) -> dict[int, str]: |
| """episode_index -> LLM-generated natural language prompt (DreamDojo).""" |
| cache = ws / "scripts/gr1_task_labels_cache.jsonl" |
| ep_to_prompt: dict[int, str] = {} |
| if not cache.exists(): |
| return ep_to_prompt |
| with cache.open() as f: |
| for line in f: |
| line = line.strip() |
| if not line: |
| continue |
| obj = json.loads(line) |
| if obj.get("source") != "DreamDojo-HV_Eval": |
| continue |
| m = re.search(r"(\d+)$", obj.get("episode", "")) |
| if m: |
| ep_to_prompt[int(m.group(1))] = obj.get("prompt", "") |
| return ep_to_prompt |
|
|
|
|
| |
| |
| |
|
|
| def sample_frames(video_path: Path, n: int) -> list[Image.Image]: |
| frames = iio.imread(str(video_path), plugin="pyav") |
| if frames.ndim == 3: |
| frames = frames[None] |
| T = frames.shape[0] |
| if T <= n: |
| idx = list(range(T)) |
| else: |
| step = T / n |
| idx = [int(i * step) for i in range(n)] |
| return [Image.fromarray(frames[i]) for i in idx] |
|
|
|
|
| def build_messages(frames: list[Image.Image], task_hint: str) -> list[dict]: |
| content = [{"type": "image", "image": img} for img in frames] |
| if task_hint: |
| prompt_text = CAPTION_PROMPT.format(task_hint=task_hint) |
| else: |
| prompt_text = CAPTION_PROMPT_NO_HINT |
| content.append({"type": "text", "text": prompt_text}) |
| return [{"role": "user", "content": content}] |
|
|
|
|
| |
| |
| |
|
|
| def parse_args(): |
| p = argparse.ArgumentParser( |
| description="Caption humanoid singleview training videos with Qwen3-VL." |
| ) |
| p.add_argument( |
| "--workspace", type=Path, default=WS, |
| help="Project root (default: auto-detected from script location).", |
| ) |
| p.add_argument( |
| "--splits", nargs="+", default=["makovian", "non_makovian"], |
| help="Which split folders under datasets/humanoid/singleview/ to process.", |
| ) |
| p.add_argument( |
| "--dreamgen_target", type=Path, |
| default=None, |
| help="Where to write humanoid_dreamgen/ flat structure. " |
| "Default: <workspace>/datasets/humanoid_dreamgen/", |
| ) |
| p.add_argument("--model", default="Qwen/Qwen3-VL-30B-A3B-Instruct-FP8") |
| p.add_argument("--num_video_frames", type=int, default=8) |
| p.add_argument("--max_new_tokens", type=int, default=80) |
| p.add_argument("--tensor_parallel_size", type=int, default=1) |
| p.add_argument("--gpu_memory_utilization", type=float, default=0.85) |
| p.add_argument("--max_model_len", type=int, default=8192) |
| p.add_argument("--batch_size", type=int, default=8) |
| p.add_argument( |
| "--overwrite", action="store_true", |
| help="Re-caption even if a .txt caption already exists.", |
| ) |
| p.add_argument( |
| "--dry_run", action="store_true", |
| help="Print what would be done without running the LLM.", |
| ) |
| return p.parse_args() |
|
|
|
|
| def main(): |
| args = parse_args() |
| ws = args.workspace.resolve() |
|
|
| sv_root = ws / "datasets/humanoid/singleview" |
| dreamgen_target = args.dreamgen_target or (ws / "datasets/humanoid_dreamgen") |
| dreamgen_videos = dreamgen_target / "videos" |
| dreamgen_metas = dreamgen_target / "metas" |
| dreamgen_videos.mkdir(parents=True, exist_ok=True) |
| dreamgen_metas.mkdir(parents=True, exist_ok=True) |
|
|
| |
| print("Loading task hint tables...") |
| gr1_hints = load_gr1robot_task_hints(ws) |
| dd_hints = load_dreamdojo_task_hints(ws) |
| print(f" GR1_robot hints: {len(gr1_hints)} episodes") |
| print(f" DreamDojo hints: {len(dd_hints)} episodes") |
|
|
| |
| |
| |
| |
| work = [] |
|
|
| for split in args.splits: |
| split_dir = sv_root / split |
| if not split_dir.exists(): |
| print(f"[warn] split dir not found: {split_dir}") |
| continue |
|
|
| videos = sorted((split_dir / "videos").rglob("*.mp4")) |
| print(f"[{split}] found {len(videos)} videos") |
|
|
| for v in videos: |
| |
| try: |
| target = os.readlink(v) |
| except OSError: |
| target = str(v) |
|
|
| if "GR1_robot" in target: |
| src_type = "GR1_robot" |
| elif "DreamDojo" in target: |
| src_type = "DreamDojo" |
| else: |
| src_type = "unknown" |
|
|
| m = re.search(r"(\d+)$", v.stem) |
| ep_id = int(m.group(1)) if m else -1 |
|
|
| if src_type == "GR1_robot": |
| task_hint = gr1_hints.get(ep_id, "") |
| elif src_type == "DreamDojo": |
| task_hint = dd_hints.get(ep_id, "") |
| else: |
| task_hint = "" |
|
|
| |
| caption_dir = split_dir / "captions" |
| caption_txt = caption_dir / f"{v.stem}.txt" |
|
|
| |
| dg_stem = f"{split}__{v.stem}" |
|
|
| work.append((split, v, ep_id, task_hint, src_type, caption_txt, dg_stem)) |
|
|
| if not work: |
| print("No videos found — nothing to do.") |
| return |
|
|
| |
| todo = [w for w in work if args.overwrite or not w[5].exists()] |
| print(f"\nTotal videos: {len(work)} | Need caption: {len(todo)}") |
|
|
| if args.dry_run: |
| print("\n[dry_run] Sample items:") |
| for w in todo[:5]: |
| split, v, ep_id, hint, src, cap_txt, dg_stem = w |
| print(f" [{split}] ep={ep_id:06d} src={src} hint='{hint[:60]}' dg_stem={dg_stem}") |
| print("... (dry_run, exiting)") |
| return |
|
|
| |
| |
| |
| if todo: |
| print(f"\nLoading {args.model} ...") |
| from vllm import LLM, SamplingParams |
| from transformers import AutoProcessor |
|
|
| processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True) |
| llm = LLM( |
| model=args.model, |
| trust_remote_code=True, |
| tensor_parallel_size=args.tensor_parallel_size, |
| gpu_memory_utilization=args.gpu_memory_utilization, |
| max_model_len=args.max_model_len, |
| limit_mm_per_prompt={"image": args.num_video_frames}, |
| dtype="auto", |
| ) |
| sampling = SamplingParams( |
| temperature=0.2, top_p=0.9, max_tokens=args.max_new_tokens |
| ) |
|
|
| B = args.batch_size |
| for i in range(0, len(todo), B): |
| chunk = todo[i : i + B] |
| prompts_input = [] |
| for split, v, ep_id, task_hint, src_type, cap_txt, dg_stem in chunk: |
| try: |
| frames = sample_frames(v, args.num_video_frames) |
| except Exception as e: |
| print(f" [warn] frame read failed {v.name}: {e}") |
| frames = [] |
| if not frames: |
| prompts_input.append(None) |
| continue |
| msgs = build_messages(frames, task_hint) |
| text = processor.apply_chat_template( |
| msgs, tokenize=False, add_generation_prompt=True |
| ) |
| prompts_input.append({ |
| "prompt": text, |
| "multi_modal_data": {"image": frames}, |
| }) |
|
|
| |
| valid = [(chunk[j], p) for j, p in enumerate(prompts_input) if p is not None] |
| if not valid: |
| continue |
|
|
| outs = llm.generate([p for _, p in valid], sampling) |
| for (split, v, ep_id, task_hint, src_type, cap_txt, dg_stem), out in zip( |
| [item for item, _ in valid], outs |
| ): |
| caption = out.outputs[0].text.strip().replace("\n", " ") |
| cap_txt.parent.mkdir(parents=True, exist_ok=True) |
| cap_txt.write_text(caption) |
|
|
| done = min(i + B, len(todo)) |
| print(f" captioned {done}/{len(todo)}") |
|
|
| |
| |
| |
| print("\nWriting humanoid_dreamgen/ flat structure...") |
| written = defaultdict(int) |
|
|
| for split, v, ep_id, task_hint, src_type, cap_txt, dg_stem in work: |
| if not cap_txt.exists(): |
| print(f" [skip] no caption for {v.stem} ({split})") |
| continue |
|
|
| caption = cap_txt.read_text().strip() |
| if not caption: |
| continue |
|
|
| |
| dg_vid = dreamgen_videos / f"{dg_stem}.mp4" |
| if not dg_vid.exists(): |
| dg_vid.symlink_to(v.resolve()) |
|
|
| |
| dg_txt = dreamgen_metas / f"{dg_stem}.txt" |
| if not dg_txt.exists() or args.overwrite: |
| dg_txt.write_text(caption) |
|
|
| written[split] += 1 |
|
|
| print("\n=== Done ===") |
| for split in args.splits: |
| n = written[split] |
| print(f" {split}: {n} episodes written to humanoid_dreamgen/") |
| print(f" Target: {dreamgen_target}") |
| print("\nNext steps:") |
| print(" 1. bash finetuning/dreamgen/scripts/compute_t5_embeddings.sh (DATASETS=humanoid)") |
| print(" 2. NPROC=8 bash finetuning/dreamgen/launch.sh predict2_video2world_training_2b_humanoid_singleview") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
