#!/usr/bin/env python3 """ Prepare InstanceV training data from iGround processed JSONL. Outputs per-line JSON with: { "video": "relative/path/to/clip.mp4", "prompt": "caption", "instance_prompts": ["phrase1", "phrase2", ...], "instance_mask_dirs": [ {"mask_dir": "/abs/path/to/masks", "instance_id": 0, "num_frames": 49}, ... ] } """ import argparse import json import math import os from pathlib import Path import imageio.v2 as imageio from PIL import Image, ImageDraw from tqdm import tqdm def parse_args(): parser = argparse.ArgumentParser(description="Prepare InstanceV data from iGround") parser.add_argument( "--iground_jsonl", type=str, default="/data/rczhang/PencilFolder/data/iGround/iGround_train_set_processed.jsonl", help="Path to iGround processed JSONL.", ) parser.add_argument( "--clips_dir", type=str, default="/data/rczhang/PencilFolder/data/iGround/Clips/train", help="Directory containing iGround clips.", ) parser.add_argument( "--mask_root_dir", type=str, default="/data/rczhang/PencilFolder/data/iGround/InstanceMasks/train", help="Root directory to store generated instance masks.", ) parser.add_argument( "--output_metadata", type=str, default="/data/rczhang/PencilFolder/data/iGround/instancev_iground_train.jsonl", help="Output metadata JSONL path.", ) parser.add_argument( "--dataset_base_path", type=str, default="/data/rczhang/PencilFolder/data", help="Base path used by UnifiedDataset (video paths will be relative to this).", ) parser.add_argument( "--min_instances", type=int, default=1, help="Minimum number of instances required.", ) parser.add_argument( "--max_instances", type=int, default=None, help="Maximum number of instances to keep (None = keep all).", ) parser.add_argument( "--overwrite_masks", action="store_true", help="Overwrite existing masks for a clip.", ) parser.add_argument( "--limit", type=int, default=None, help="Limit number of samples for debugging.", ) return parser.parse_args() def _safe_relpath(path: str, base_path: str) -> str: if not base_path: return path return os.path.relpath(path, base_path) def _clamp_bbox(bbox, width: int, height: int): if not bbox or len(bbox) != 4: return None x0, y0, x1, y1 = bbox left = max(0, int(math.floor(x0))) top = max(0, int(math.floor(y0))) right = min(width, int(math.ceil(x1))) bottom = min(height, int(math.ceil(y1))) if right <= left or bottom <= top: return None return left, top, right, bottom def _collect_visible_phrases(phrases, labels_per_frame): visible = set() for labels in labels_per_frame: for label in labels: visible.add(label) return [p for p in phrases if p in visible] def _write_masks( mask_dir: str, phrases, labels_per_frame, bboxes_per_frame, width: int, height: int, overwrite: bool, ): if os.path.isdir(mask_dir) and not overwrite: return os.makedirs(mask_dir, exist_ok=True) phrase_set = set(phrases) num_frames = len(bboxes_per_frame) for frame_idx in range(num_frames): labels = labels_per_frame[frame_idx] bboxes = bboxes_per_frame[frame_idx] frame_map = {} for label, bbox in zip(labels, bboxes): if label in phrase_set: frame_map[label] = bbox for inst_id, phrase in enumerate(phrases): mask = Image.new("L", (width, height), 0) bbox = frame_map.get(phrase) if bbox is not None: coords = _clamp_bbox(bbox, width, height) if coords is not None: draw = ImageDraw.Draw(mask) draw.rectangle(coords, fill=255) mask_path = os.path.join(mask_dir, f"{frame_idx:06d}_No.{inst_id}.png") mask.save(mask_path) def _is_video_readable(video_path: str) -> bool: try: reader = imageio.get_reader(video_path) try: reader.get_data(0) finally: reader.close() except Exception: return False return True def main(): args = parse_args() Path(args.mask_root_dir).mkdir(parents=True, exist_ok=True) Path(os.path.dirname(args.output_metadata)).mkdir(parents=True, exist_ok=True) processed = 0 skipped_missing_video = 0 skipped_instances = 0 skipped_unreadable = 0 wrote = 0 with open(args.iground_jsonl, "r", encoding="utf-8") as f_in, open( args.output_metadata, "w", encoding="utf-8" ) as f_out: for line in tqdm(f_in, desc="Processing iGround"): if args.limit is not None and wrote >= args.limit: break line = line.strip() if not line: continue processed += 1 sample = json.loads(line) video_id = sample["video_id"] clip_id = sample["clip_id"] clip_name = f"{video_id}_{clip_id}.mp4" clip_path = os.path.join(args.clips_dir, clip_name) if not os.path.isfile(clip_path): skipped_missing_video += 1 continue if not _is_video_readable(clip_path): skipped_unreadable += 1 continue phrases = list(sample.get("phrases", [])) labels_per_frame = sample.get("labels", []) bboxes_per_frame = sample.get("bboxes", []) if not phrases or not labels_per_frame or not bboxes_per_frame: skipped_instances += 1 continue visible_phrases = _collect_visible_phrases(phrases, labels_per_frame) if args.max_instances is not None: visible_phrases = visible_phrases[: args.max_instances] if len(visible_phrases) < args.min_instances: skipped_instances += 1 continue width = int(sample["width"]) height = int(sample["height"]) mask_dir = os.path.join(args.mask_root_dir, f"{video_id}_{clip_id}_masks") _write_masks( mask_dir, visible_phrases, labels_per_frame, bboxes_per_frame, width, height, overwrite=args.overwrite_masks, ) instance_mask_dirs = [ { "mask_dir": mask_dir, "instance_id": inst_id, "num_frames": len(bboxes_per_frame), } for inst_id in range(len(visible_phrases)) ] entry = { "video": _safe_relpath(clip_path, args.dataset_base_path), "prompt": sample.get("caption", ""), "instance_prompts": visible_phrases, "instance_mask_dirs": instance_mask_dirs, } f_out.write(json.dumps(entry, ensure_ascii=False) + "\n") wrote += 1 print("Done.") print(f"Processed: {processed}") print(f"Wrote: {wrote}") print(f"Skipped (missing video): {skipped_missing_video}") print(f"Skipped (unreadable video): {skipped_unreadable}") print(f"Skipped (insufficient instances): {skipped_instances}") if __name__ == "__main__": main()