130 lines
5.3 KiB
Python
130 lines
5.3 KiB
Python
"""Convert dataset captions into LingBot-Video structured JSON captions.
|
|
|
|
The LingBot-Video DiT was trained ONLY on structured JSON captions produced by
|
|
the official prompt rewriter. Training with plain natural-language captions
|
|
feeds out-of-distribution text to the model and degrades it. Run this script
|
|
BEFORE training to rewrite the "text" field of the dataset metadata:
|
|
|
|
# needs the official rewriter weights:
|
|
# REWRITER_BASE_MODEL -> Qwen3.6-27B VLM (or --base)
|
|
# REWRITER_ADAPTER -> step2 LoRA adapter, peft format (or --adapter)
|
|
python scripts/lingbot_video/prepare_captions.py \
|
|
--metadata datasets/my_dataset/metadata.json \
|
|
--data_root datasets/my_dataset \
|
|
--output datasets/my_dataset/metadata_json.json \
|
|
--mode ti2v --duration 3.3
|
|
|
|
Then point train.sh at the rewritten metadata
|
|
(DATASET_META_NAME=datasets/my_dataset/metadata_json.json).
|
|
|
|
Behaviour:
|
|
- entries whose "text" is already a valid structured JSON caption are kept
|
|
as-is (use --overwrite to re-rewrite them);
|
|
- rewriting is resumed safely: the output file is saved after every sample;
|
|
- for --mode ti2v the video's first frame is fed to the rewriter;
|
|
- the original metadata file is never modified in place.
|
|
"""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
|
|
current_file_path = os.path.abspath(__file__)
|
|
project_root = os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))
|
|
sys.path.insert(0, project_root)
|
|
|
|
from videox_fun.models.lingbot_video_rewriter import is_valid_caption, LingBotVideoRewriter
|
|
|
|
|
|
def load_metadata(path):
|
|
if path.endswith(".csv"):
|
|
import csv
|
|
with open(path, "r") as f:
|
|
return list(csv.DictReader(f))
|
|
with open(path, encoding="utf-8") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def save_metadata(entries, path):
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump(entries, f, ensure_ascii=False, indent=2)
|
|
|
|
|
|
def extract_first_frame(video_path):
|
|
"""Return the first frame of a video as a PIL image (ti2v condition)."""
|
|
from PIL import Image
|
|
try:
|
|
from decord import VideoReader
|
|
vr = VideoReader(video_path)
|
|
return Image.fromarray(vr[0].asnumpy())
|
|
except ImportError:
|
|
import cv2
|
|
cap = cv2.VideoCapture(video_path)
|
|
ok, frame = cap.read()
|
|
cap.release()
|
|
if not ok:
|
|
raise RuntimeError(f"cannot read first frame of {video_path}")
|
|
return Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--metadata", required=True, help="dataset metadata (.json list or .csv)")
|
|
ap.add_argument("--data_root", default="", help="dataset root used to resolve file_path")
|
|
ap.add_argument("--output", required=True, help="path of the rewritten metadata")
|
|
ap.add_argument("--mode", default="ti2v", choices=["t2v", "ti2v", "t2i"])
|
|
ap.add_argument("--duration", type=float, default=5.0,
|
|
help="clip duration in seconds fed to the rewriter (match your training clip)")
|
|
ap.add_argument("--base", default=os.environ.get("REWRITER_BASE_MODEL", ""),
|
|
help="rewriter base VLM path (or set REWRITER_BASE_MODEL)")
|
|
ap.add_argument("--adapter", default=os.environ.get("REWRITER_ADAPTER", ""),
|
|
help="rewriter step2 LoRA adapter path (or set REWRITER_ADAPTER)")
|
|
ap.add_argument("--overwrite", action="store_true",
|
|
help="re-rewrite captions that are already valid JSON captions")
|
|
args = ap.parse_args()
|
|
|
|
entries = load_metadata(args.metadata)
|
|
print(f"loaded {len(entries)} entries from {args.metadata}")
|
|
|
|
rewriter = LingBotVideoRewriter(base=args.base, adapter=args.adapter)
|
|
|
|
n_done = n_kept = n_fail = 0
|
|
for i, entry in enumerate(entries):
|
|
text = entry.get("text", "")
|
|
if is_valid_caption(text) and not args.overwrite:
|
|
n_kept += 1
|
|
continue
|
|
file_path = entry.get("file_path", "")
|
|
first_frame = None
|
|
if args.mode == "ti2v" and file_path:
|
|
video_path = file_path if os.path.isabs(file_path) else os.path.join(args.data_root, file_path)
|
|
try:
|
|
first_frame = extract_first_frame(video_path)
|
|
except Exception as e:
|
|
print(f"[{i}] WARN: cannot read first frame of {video_path}: {e}")
|
|
try:
|
|
entry["text"] = rewriter.rewrite(text, args.mode, first_frame, args.duration)
|
|
caption = entry["text"]
|
|
except Exception as e:
|
|
caption = None
|
|
print(f"[{i}] ERROR rewriting {file_path}: {e}")
|
|
if caption is None:
|
|
n_fail += 1
|
|
print(f"[{i}] FAILED, keep original text: {file_path}")
|
|
else:
|
|
n_done += 1
|
|
# checkpoint after every sample so long runs can be interrupted
|
|
save_metadata(entries, args.output)
|
|
print(f"[{i + 1}/{len(entries)}] rewritten={n_done} kept={n_kept} failed={n_fail} -> {file_path}")
|
|
|
|
save_metadata(entries, args.output)
|
|
print(f"\nDONE: rewritten={n_done}, already-valid={n_kept}, failed={n_fail}\noutput: {args.output}")
|
|
if n_fail:
|
|
print("NOTE: failed entries still carry their original natural-language text; "
|
|
"fix them manually (see videox_fun/pipeline/lingbot_video_caption.py) "
|
|
"or re-run this script.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|