716 lines
40 KiB
Python
716 lines
40 KiB
Python
"""Extract B1 frozen Qwen local crops + recent mixed audio, aligned to Omni .pt.
|
|
|
|
Uses provided boxes (not speaking labels), raw source media and the existing
|
|
dense-timeline reader. No Thinker forward, cross-person attention, or learned
|
|
256D projection is performed here. See docs/local_feature_baseline.md.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
from tqdm.auto import tqdm
|
|
|
|
from easycom_head import load_ego4d_manifest_split, normalize_sample_name, resolve_feature_path
|
|
from local_turn_features import (
|
|
LOCAL_FORMAT, causal_frame_ticks, crop_box, easycom_boxes,
|
|
load_local_features, pooled_recent_visual, reference_alignment,
|
|
)
|
|
|
|
|
|
TESTED_TRANSFORMERS_VERSIONS = ("4.52.3", "4.57.6")
|
|
|
|
|
|
def validate_local_transformers_version(version: str) -> str:
|
|
"""Keep the local encoder's tested API separate from the replay KV API."""
|
|
if version not in TESTED_TRANSFORMERS_VERSIONS:
|
|
supported = ", ".join(TESTED_TRANSFORMERS_VERSIONS)
|
|
raise RuntimeError(
|
|
f"Local encoder API is tested with transformers versions {supported}; found {version}. "
|
|
"Use a tested version or validate the encoder APIs before adding support."
|
|
)
|
|
return version
|
|
|
|
|
|
def model_config_digest(config: dict[str, Any]) -> str:
|
|
"""Ignore location fields so an explicitly identified checkpoint can move."""
|
|
def stable(value):
|
|
if isinstance(value, dict):
|
|
return {key: stable(item) for key, item in value.items() if key != "_name_or_path"}
|
|
if isinstance(value, (list, tuple)):
|
|
return [stable(item) for item in value]
|
|
return value
|
|
|
|
return hashlib.sha256(json.dumps(stable(config), sort_keys=True).encode()).hexdigest()
|
|
|
|
|
|
class FrozenQwenLocalEncoder:
|
|
"""Keep only frozen vision+merger and audio+projection on the GPU."""
|
|
|
|
def __init__(self, args: argparse.Namespace) -> None:
|
|
import transformers
|
|
from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniThinkerForConditionalGeneration
|
|
|
|
self.transformers_version = validate_local_transformers_version(transformers.__version__)
|
|
print(f"Local encoders: transformers={self.transformers_version}", flush=True)
|
|
self.device = torch.device(args.device)
|
|
if self.device.type == "cuda" and not torch.cuda.is_available():
|
|
raise RuntimeError(f"CUDA requested but unavailable: {args.device}")
|
|
dtype_key = "dtype" if self.transformers_version == "4.57.6" else "torch_dtype"
|
|
kwargs = {dtype_key: getattr(torch, args.dtype), "device_map": {"": "cpu"},
|
|
"output_loading_info": True}
|
|
if args.flash_attention_2:
|
|
kwargs["attn_implementation"] = "flash_attention_2"
|
|
model, info = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained(args.model_id, **kwargs)
|
|
bad = [key for key in info.get("missing_keys", [])
|
|
if key.startswith(("visual.", "audio_tower."))]
|
|
if bad or info.get("mismatched_keys") or info.get("error_msgs"):
|
|
raise RuntimeError(f"Cannot extract frozen features from incompletely loaded weights: {info}")
|
|
model.requires_grad_(False).eval()
|
|
self.visual_dim = int(model.config.vision_config.out_hidden_size)
|
|
self.audio_dim = int(model.config.audio_config.output_dim)
|
|
self.model_config_sha256 = model_config_digest(model.config.to_dict())
|
|
# These modules are never called. Releasing them avoids keeping the
|
|
# multi-billion-parameter language decoder on the extraction GPU.
|
|
model.model = None
|
|
model.lm_head = None
|
|
model.visual.to(self.device)
|
|
model.audio_tower.to(self.device)
|
|
self.model = model
|
|
# Keep PIL/slow image preprocessing consistent with the original B1
|
|
# definition, independently of newer Transformers' fast defaults.
|
|
self.processor = Qwen2_5OmniProcessor.from_pretrained(args.model_id, use_fast=False)
|
|
self.max_pixels = args.max_pixels
|
|
self.image_batch_size = args.image_batch_size
|
|
|
|
@torch.inference_mode()
|
|
def images(self, crops: list[Any]) -> torch.Tensor:
|
|
pieces = []
|
|
for start in range(0, len(crops), self.image_batch_size):
|
|
batch = self.processor.image_processor(
|
|
images=crops[start:start + self.image_batch_size],
|
|
min_pixels=4 * 28 * 28, max_pixels=self.max_pixels, return_tensors="pt",
|
|
)
|
|
grid = batch["image_grid_thw"].to(self.device)
|
|
pixels = batch["pixel_values"].to(self.device, dtype=self.model.visual.dtype)
|
|
tokens = self.model.get_image_features(pixels, grid)
|
|
merge = int(self.model.config.vision_config.spatial_merge_size)
|
|
counts = (grid.prod(-1) // (merge * merge)).tolist()
|
|
if not isinstance(tokens, torch.Tensor) or tokens.shape != (sum(counts), self.visual_dim):
|
|
raise RuntimeError(f"Unexpected image encoder output: {type(tokens).__name__}, "
|
|
f"shape={getattr(tokens, 'shape', None)}")
|
|
# Split BEFORE spatial mean: each person has one vote regardless of
|
|
# box size / number of visual tokens.
|
|
pieces.extend(part.float().mean(0).cpu() for part in tokens.split(counts))
|
|
return torch.stack(pieces) if pieces else torch.empty(0, self.visual_dim)
|
|
|
|
@torch.inference_mode()
|
|
def audio(self, waveform: np.ndarray) -> torch.Tensor:
|
|
if waveform.ndim != 1 or not len(waveform) or not np.isfinite(waveform).all():
|
|
raise ValueError("Audio must be finite nonempty mono samples")
|
|
# The waveform has ALREADY been cut at the current time. Computing its
|
|
# mel transform here cannot read future samples of the source file.
|
|
batch = self.processor.feature_extractor(
|
|
waveform, sampling_rate=16000, padding="longest",
|
|
return_attention_mask=True, return_tensors="pt",
|
|
)
|
|
features = batch["input_features"].to(self.device, dtype=self.model.audio_tower.dtype)
|
|
mask = batch["attention_mask"].to(self.device)
|
|
tokens = self.model.get_audio_features(features, feature_attention_mask=mask)
|
|
_, lengths = self.model.audio_tower._get_feat_extract_output_lengths(mask.sum(-1))
|
|
if (not isinstance(tokens, torch.Tensor)
|
|
or tokens.shape != (int(lengths.sum()), self.audio_dim) or not len(tokens)):
|
|
raise RuntimeError(f"Unexpected valid audio encoder output: {type(tokens).__name__}, "
|
|
f"shape={getattr(tokens, 'shape', None)}")
|
|
# get_audio_features removes padding and verifies the output length.
|
|
return tokens.float().mean(0).cpu()
|
|
|
|
|
|
def ego4d_clip_boxes(clip: dict[str, Any]) -> dict[int, list[dict[str, Any]]]:
|
|
"""AVD tracks indexed by clip-local decoded frame; never use voice labels."""
|
|
result: dict[int, list[dict[str, Any]]] = {}
|
|
for person in clip.get("persons", []):
|
|
if person.get("camera_wearer", False):
|
|
continue
|
|
seen = set()
|
|
for track in person.get("tracking_paths", []):
|
|
for box in track.get("track", []):
|
|
frame = box.get("clip_frame")
|
|
if frame is None:
|
|
frame = box.get("frame")
|
|
if frame is None:
|
|
frame = int(box["video_frame"]) - int(clip["video_start_frame"])
|
|
frame = int(frame)
|
|
if frame < 0:
|
|
raise ValueError("Negative clip-relative Ego4D box frame")
|
|
# Some exports duplicate a person's boundary box in two tracks.
|
|
if frame in seen:
|
|
continue
|
|
seen.add(frame)
|
|
x, y, w, h = (float(box[k]) for k in ("x", "y", "width", "height"))
|
|
result.setdefault(frame, []).append(
|
|
{"id": str(person["person_id"]), "xyxy": [x, y, x + w, y + h]}
|
|
)
|
|
return result
|
|
|
|
|
|
def load_ego4d_boxes(paths: list[Path], names: list[str]) -> dict[str, Any]:
|
|
result = {}
|
|
selected = set(names)
|
|
for path in paths:
|
|
print(f"Reading provided tracks: {path}", flush=True)
|
|
with path.open() as handle:
|
|
payload = json.load(handle)
|
|
for video in payload["videos"]:
|
|
for clip in video["clips"]:
|
|
uid = clip["clip_uid"]
|
|
if uid in selected:
|
|
if uid in result:
|
|
raise ValueError(f"Duplicate Ego4D annotation UID: {uid}")
|
|
# Retain declared bounds so a terminal record can be
|
|
# distinguished from a truncated video or frame mismatch.
|
|
result[uid] = {
|
|
"frames": ego4d_clip_boxes(clip),
|
|
"bounds": {key: clip.get(key) for key in (
|
|
"clip_start_frame", "clip_end_frame", "clip_start_sec", "clip_end_sec",
|
|
)},
|
|
}
|
|
del payload
|
|
missing = set(names) - result.keys()
|
|
if missing:
|
|
raise ValueError(f"No provided tracks for: {sorted(missing)}")
|
|
return result
|
|
|
|
|
|
def validate_ego4d_box_frames(name: str, annotation: dict[str, Any], source):
|
|
"""Ignore only a confirmed exclusive-end record; never shift/clamp boxes."""
|
|
frames, bounds = annotation["frames"], annotation["bounds"]
|
|
count = source.num_frames
|
|
outside = sorted(frame for frame in frames if not 0 <= frame < count)
|
|
audit = {
|
|
"policy": "clip-local-zero-based; ignore-confirmed-exclusive-end-only",
|
|
"source_frames": count, "source_fps": source.source_fps,
|
|
"declared_bounds": dict(bounds),
|
|
"annotation_frame_range": [min(frames), max(frames)] if frames else None,
|
|
"ignored_frame_indices": [], "ignored_box_count": 0,
|
|
}
|
|
if not outside:
|
|
return frames, audit
|
|
start_sec, end_sec = bounds.get("clip_start_sec"), bounds.get("clip_end_sec")
|
|
duration_matches = (
|
|
isinstance(start_sec, (int, float)) and isinstance(end_sec, (int, float))
|
|
and np.isfinite([start_sec, end_sec, source.source_fps]).all()
|
|
and source.source_fps > 0 and abs(start_sec) <= 1e-6
|
|
and abs((end_sec - start_sec) - count / source.source_fps) <= 1e-4
|
|
)
|
|
if (outside != [count] or bounds.get("clip_start_frame") != 0
|
|
or bounds.get("clip_end_frame") != count or not duration_matches):
|
|
raise ValueError(
|
|
f"Ego4D annotation frame exceeds source clip: {name}; "
|
|
f"source_frames={count}, source_fps={source.source_fps}, "
|
|
f"annotation_range={audit['annotation_frame_range']}, "
|
|
f"out_of_range_count={len(outside)}, first_10={outside[:10]}, "
|
|
f"declared_bounds={bounds}. Only frame == source_frames == clip_end_frame "
|
|
"with zero-origin, matching duration is an allowed terminal record; "
|
|
"check annotation/video alignment instead of clamping or skipping the clip."
|
|
)
|
|
audit.update(ignored_frame_indices=[count], ignored_box_count=len(frames[count]))
|
|
print(f"{name}: ignoring exclusive-end bbox record at frame {count} "
|
|
f"({len(frames[count])} boxes); decoded frames are 0..{count - 1}. "
|
|
"No frame shift or feature/target row removal.", flush=True)
|
|
return {frame: boxes for frame, boxes in frames.items() if frame != count}, audit
|
|
|
|
|
|
def build_media_and_boxes(args: argparse.Namespace, name: str, end: float, ego_boxes: dict[str, Any]):
|
|
# Lazy imports allow --help and pure tests without Transformers/decoders.
|
|
from streaming_turn_model_directly import CausalSourceAV, SourceClipSpec, _easycom_directory_specs
|
|
|
|
if args.dataset == "easycom":
|
|
specs = _easycom_directory_specs(args.data_root / "Main" / "Video_Compressed" / name)
|
|
else:
|
|
specs = [SourceClipSpec(args.data_root / "v1" / "clips" / f"{name}.mp4")]
|
|
media = CausalSourceAV(specs, output_fps=args.video_fps, required_duration=end,
|
|
decoder_backend=args.decoder_backend, decoder_threads=args.decoder_threads,
|
|
decoder_gpu_id=args.decoder_gpu_id)
|
|
if args.dataset == "easycom":
|
|
folder = "Face_Bounding_Boxes" if args.box_kind == "face" else "Head_Bounding_Boxes"
|
|
annotations = []
|
|
for clip in media.clips:
|
|
path = args.data_root / "Main" / folder / name / f"{clip.path.stem}.json"
|
|
with path.open() as handle:
|
|
rows = json.load(handle)
|
|
mapping = easycom_boxes(rows, args.box_kind)
|
|
# Head annotations have a documented extra terminal record; it is
|
|
# not a decoded frame and is never sampled.
|
|
allowed = clip.num_frames + (1 if args.box_kind == "head" else 0)
|
|
if set(mapping) != set(range(allowed)):
|
|
raise ValueError(f"Box/video frame inventory mismatch: {path}")
|
|
annotations.append(mapping)
|
|
else:
|
|
mapping, audit = validate_ego4d_box_frames(name, ego_boxes[name], media.clips[0])
|
|
annotations = [mapping]
|
|
ego_boxes[name]["frame_audit"] = audit
|
|
return media, annotations
|
|
|
|
|
|
def box_coordinate_offset(clip, coordinate_space: str) -> tuple[int, int]:
|
|
"""Map source/encoded boxes into the image actually returned by the reader."""
|
|
if coordinate_space == "encoded":
|
|
return tuple(clip.crop[:2]) if clip.crop is not None else (0, 0)
|
|
if coordinate_space == "camera":
|
|
# Only for annotations already translated into the cropped image.
|
|
return (0, 0)
|
|
raise ValueError(f"Unknown box coordinate space: {coordinate_space}")
|
|
|
|
|
|
def frame_crops(media, annotations, time: float, args: argparse.Namespace):
|
|
clip_index, frame_index, source_pts = media._source_frame_for_time(time)
|
|
clip = media.clips[clip_index]
|
|
image = media.frame_at_or_before(time)
|
|
# EasyCom's supplied boxes refer to the encoded image, including the left
|
|
# banner. The reader has already removed that banner. Subtract this clip's
|
|
# actual crop origin, not a hard-coded offset shared by all sessions.
|
|
offset_x, offset_y = box_coordinate_offset(clip, args.box_coordinate_space)
|
|
crops, records = [], []
|
|
for person in annotations[clip_index].get(frame_index, []):
|
|
box = crop_box(person["xyxy"], *image.size, offset_x=offset_x, offset_y=offset_y,
|
|
expansion=args.crop_expansion)
|
|
if box is not None:
|
|
crops.append(image.crop(box))
|
|
records.append({"id": person["id"], "box": box,
|
|
"annotation_box": list(person["xyxy"]),
|
|
"coordinate_offset": [offset_x, offset_y],
|
|
"source_clip_index": clip_index,
|
|
"source_frame_index": frame_index,
|
|
"source_time": clip.dense_start + source_pts})
|
|
return image, crops, records
|
|
|
|
|
|
def validate_source_timeline(reference: dict[str, Any], media) -> None:
|
|
"""Check source order/frames/offsets when the Omni file has direct metadata."""
|
|
saved = (reference.get("direct_media") or {}).get("source_clips")
|
|
if saved is None:
|
|
print("NOTE: reference has no direct-media inventory; verify its dense source timeline manually.")
|
|
return
|
|
if len(saved) < len(media.clips):
|
|
raise ValueError("Local source has more clips than the reference inventory")
|
|
for old, new in zip(saved, media.metadata()["source_clips"]):
|
|
if (Path(old["path"]).name != Path(new["path"]).name
|
|
or old["source_frames"] != new["source_frames"]
|
|
or abs(old["dense_start"] - new["dense_start"]) > 1e-5
|
|
or abs(old["duration"] - new["duration"]) > 1e-5):
|
|
raise ValueError(f"Local/reference source timeline differs at {new['path']}")
|
|
|
|
|
|
def preview_time_bounds(args: argparse.Namespace, available_end: float) -> tuple[float, float]:
|
|
"""Preview-only range on the original timeline, never a rebased feature grid."""
|
|
start = float(getattr(args, "preview_start", 0.))
|
|
requested_end = getattr(args, "preview_end", None)
|
|
end = min(available_end, requested_end) if requested_end is not None else available_end
|
|
if start >= end:
|
|
raise ValueError(f"Preview start {start:g}s is outside the available prefix ending at {end:g}s; "
|
|
"remove/increase --max-steps or choose an earlier --preview-start")
|
|
return start, end
|
|
|
|
|
|
def preview_ticks(media, annotations, args: argparse.Namespace, end: float) -> tuple[list[int], int, int]:
|
|
start, stop = preview_time_bounds(args, end)
|
|
ticks = [tick for tick in causal_frame_ticks(stop, args.video_fps)
|
|
if tick / args.video_fps >= start - 1e-8]
|
|
grid_count = len(ticks)
|
|
if getattr(args, "preview_with_boxes", False):
|
|
eligible = []
|
|
# Scan annotation metadata, not every RGB frame. Use the same source
|
|
# frame mapping, coordinate correction and crop validity as extraction.
|
|
for tick in ticks:
|
|
clip_index, frame_index, _ = media._source_frame_for_time(tick / args.video_fps)
|
|
clip = media.clips[clip_index]
|
|
width, height = clip.crop[2:] if clip.crop is not None else (clip.width, clip.height)
|
|
offset_x, offset_y = box_coordinate_offset(clip, args.box_coordinate_space)
|
|
if any(crop_box(person["xyxy"], width, height, offset_x=offset_x, offset_y=offset_y,
|
|
expansion=args.crop_expansion) is not None
|
|
for person in annotations[clip_index].get(frame_index, [])):
|
|
eligible.append(tick)
|
|
ticks = eligible
|
|
eligible_count = len(ticks)
|
|
count = min(getattr(args, "preview_count", 4), eligible_count)
|
|
indices = np.linspace(0, eligible_count - 1, count, dtype=int) if count else []
|
|
return [ticks[index] for index in indices], grid_count, eligible_count
|
|
|
|
|
|
def preview_boxes(media, annotations, args: argparse.Namespace, name: str, end: float) -> None:
|
|
from PIL import ImageDraw
|
|
|
|
folder = args.output_root / f"{name}_preview"
|
|
folder.mkdir(parents=True, exist_ok=True)
|
|
start, stop = preview_time_bounds(args, end)
|
|
ticks, grid_count, eligible_count = preview_ticks(media, annotations, args, end)
|
|
print(f"Preview [{start:g},{stop:g})s: {grid_count} grid frames, {eligible_count} eligible, "
|
|
f"saving {len(ticks)} images", flush=True)
|
|
if not ticks:
|
|
print("No eligible preview frames. Try a longer range or another clip. "
|
|
"If people are visible but no valid boxes exist, check annotation/frame alignment.", flush=True)
|
|
audit = []
|
|
for tick in ticks:
|
|
time = tick / args.video_fps
|
|
image, _, records = frame_crops(media, annotations, time, args)
|
|
clip_index, frame_index, _ = media._source_frame_for_time(time)
|
|
clip = media.clips[clip_index]
|
|
offset = box_coordinate_offset(clip, args.box_coordinate_space)
|
|
draw = ImageDraw.Draw(image)
|
|
for record in records:
|
|
draw.rectangle(record["box"], outline="lime", width=4)
|
|
draw.text(record["box"][:2], record["id"], fill="red")
|
|
caption = f"t={time:.2f}s | source frame={frame_index} | people={len(records)}"
|
|
draw.rectangle(draw.textbbox((8, 8), caption), fill="black")
|
|
draw.text((8, 8), caption, fill="white")
|
|
filename = f"frame_{tick:06d}.jpg"
|
|
image.save(folder / filename)
|
|
audit.append({
|
|
"image": filename, "time": time, "source_clip": str(clip.path),
|
|
"source_frame_index": frame_index, "source_size": [clip.width, clip.height],
|
|
"crop": list(clip.crop) if clip.crop is not None else None,
|
|
"box_coordinate_space": args.box_coordinate_space,
|
|
"coordinate_offset": list(offset), "boxes": records,
|
|
})
|
|
print(f"preview t={time:g}s: {clip.width}x{clip.height} -> {image.width}x{image.height}; "
|
|
f"boxes={args.box_coordinate_space}, subtract offset={offset}, people={len(records)}")
|
|
(folder / "boxes.json").write_text(json.dumps(audit, indent=2) + "\n")
|
|
(folder / "preview.json").write_text(json.dumps({
|
|
"start": start, "end": stop, "video_fps": args.video_fps,
|
|
"with_boxes_only": getattr(args, "preview_with_boxes", False),
|
|
"requested_count": getattr(args, "preview_count", 4),
|
|
"grid_frame_count": grid_count, "eligible_frame_count": eligible_count,
|
|
"saved_count": len(audit), "saved_with_boxes": sum(bool(row["boxes"]) for row in audit),
|
|
"images": [row["image"] for row in audit],
|
|
}, indent=2) + "\n")
|
|
print(f"Box previews: {folder} (inspect before extracting all sessions)")
|
|
print("boxes.json lists this run's images; older images in the directory are not removed.")
|
|
|
|
|
|
def extract_local(media, annotations, encoder, times: torch.Tensor, args: argparse.Namespace) -> dict[str, Any]:
|
|
# Canonical grid, not float32-serialized rounding near a tick boundary.
|
|
ends = (torch.arange(len(times), dtype=torch.float64) + 1) / 5
|
|
ticks = list(causal_frame_ticks(float(ends[-1]), args.video_fps))
|
|
frame_times = torch.tensor(ticks, dtype=torch.float64) / args.video_fps
|
|
frame_vectors = torch.zeros(len(ticks), encoder.visual_dim)
|
|
frame_valid = torch.zeros(len(ticks), dtype=torch.bool)
|
|
people, records, offsets = [], [], [0]
|
|
for row, tick in enumerate(tqdm(ticks, desc="local vision", unit="frame", disable=args.no_progress)):
|
|
_, crops, observations = frame_crops(media, annotations, tick / args.video_fps, args)
|
|
vectors = encoder.images(crops)
|
|
if vectors.shape != (len(crops), encoder.visual_dim) or not torch.isfinite(vectors).all():
|
|
raise RuntimeError("Invalid person features")
|
|
if len(vectors):
|
|
stored = vectors.to(torch.bfloat16)
|
|
frame_vectors[row] = stored.float().mean(0)
|
|
frame_valid[row] = True
|
|
people.append(stored)
|
|
records.extend(observations)
|
|
offsets.append(offsets[-1] + len(vectors))
|
|
visual = torch.zeros(len(times), encoder.visual_dim)
|
|
audio = torch.zeros(len(times), encoder.audio_dim)
|
|
visual_valid = torch.zeros(len(times), dtype=torch.bool)
|
|
audio_valid = torch.zeros(len(times), dtype=torch.bool)
|
|
for row, end in enumerate(tqdm(ends.tolist(), desc="local audio/pool", unit="state", disable=args.no_progress)):
|
|
visual[row], present = pooled_recent_visual(
|
|
frame_times, frame_vectors, frame_valid, end, args.visual_window_seconds,
|
|
)
|
|
visual_valid[row] = present
|
|
waveform = media.audio_slice(max(0., end - args.audio_window_seconds), end)
|
|
audio[row] = encoder.audio(waveform)
|
|
audio_valid[row] = True # Silence is an observation, not missing audio.
|
|
if not torch.isfinite(audio).all() or not torch.isfinite(visual).all():
|
|
raise RuntimeError("Non-finite local encoder outputs")
|
|
return {
|
|
"visual": visual.to(torch.bfloat16), "audio": audio.to(torch.bfloat16),
|
|
"visual_valid": visual_valid, "audio_valid": audio_valid,
|
|
"persons": {"features": torch.cat(people) if people else torch.empty(0, encoder.visual_dim, dtype=torch.bfloat16),
|
|
"frame_times": frame_times, "frame_offsets": torch.tensor(offsets, dtype=torch.long),
|
|
"records": records},
|
|
}
|
|
|
|
|
|
def select_samples(args: argparse.Namespace) -> list[str]:
|
|
if args.manifest is not None:
|
|
names = load_ego4d_manifest_split(args.manifest, args.split)
|
|
# One-based inclusive range in the selected split, then an optional limit.
|
|
names = names[args.start_index - 1:args.end_index]
|
|
if args.limit is not None:
|
|
names = names[:args.limit]
|
|
else:
|
|
names = [normalize_sample_name(value, args.dataset) for value in args.samples]
|
|
if not names:
|
|
raise ValueError("No samples selected")
|
|
if len(set(names)) != len(names):
|
|
raise ValueError("Duplicate samples")
|
|
num_shards = getattr(args, "num_shards", 1)
|
|
shard_index = getattr(args, "shard_index", 0)
|
|
names = names[shard_index::num_shards]
|
|
return names
|
|
|
|
|
|
def requested_local_config(args: argparse.Namespace) -> dict[str, Any]:
|
|
"""Settings that must match before reusing an existing complete sidecar."""
|
|
return {
|
|
"encoder_id": args.encoder_id or args.model_id,
|
|
"dtype": args.dtype,
|
|
"attention": "flash_attention_2" if args.flash_attention_2 else "default",
|
|
"visual_fps": args.video_fps, "visual_window_seconds": args.visual_window_seconds,
|
|
"audio_window_seconds": args.audio_window_seconds,
|
|
"visual_pooling": "spatial-mean_then_equal-person-mean_then_equal-valid-frame-mean",
|
|
"audio_pooling": "valid-output-token-mean", "audio_source": "video-mixed-mono-16khz",
|
|
"box_kind": args.box_kind if args.dataset == "easycom" else "ego4d-provided-tracks",
|
|
"box_coordinate_space": args.box_coordinate_space, "crop_expansion": args.crop_expansion,
|
|
"max_pixels": args.max_pixels, "uses_speech_labels": False,
|
|
"causal_policy": "half-open-trailing-window", "box_protocol": "provided-boxes-not-online-detector",
|
|
}
|
|
|
|
|
|
def validate_reusable_sidecar(args, name, reference, version, expected_contract=None):
|
|
path = args.output_root / f"{name}.pt"
|
|
_, contract = load_local_features(path, reference, sample_id=name, dataset=args.dataset,
|
|
modalities="av", expected_contract=expected_contract)
|
|
config = contract["config"]
|
|
expected = requested_local_config(args)
|
|
if version is not None:
|
|
expected["transformers"] = version
|
|
mismatches = [key for key, value in expected.items() if config.get(key) != value]
|
|
if mismatches or not config.get("model_config_sha256") or not config.get("transformers"):
|
|
raise ValueError(f"Cannot reuse {path}: incompatible/incomplete encoder settings {mismatches}; "
|
|
"use a new output root or --overwrite")
|
|
return contract
|
|
|
|
|
|
def parse_args(argv=None):
|
|
p = argparse.ArgumentParser(description=__doc__)
|
|
p.add_argument("--dataset", choices=("easycom", "ego4d"), default="easycom")
|
|
p.add_argument("--data-root", required=True, type=Path)
|
|
p.add_argument("--feature-root", required=True, type=Path,
|
|
help="Flat .pt directory or Ego4D parent containing train/ and test/")
|
|
p.add_argument("--target-root", required=True, type=Path)
|
|
p.add_argument("--output-root", required=True, type=Path)
|
|
selection = p.add_mutually_exclusive_group(required=True)
|
|
selection.add_argument("--samples", nargs="+", help="EasyCom session numbers or Ego4D clip UIDs")
|
|
selection.add_argument("--manifest", type=Path, help="Ego4D EgoSpeak TSV; never creates a new split")
|
|
p.add_argument("--split", choices=("all", "train", "test"), default="all")
|
|
p.add_argument("--start-index", type=int, default=1, help="One-based index AFTER split filtering")
|
|
p.add_argument("--end-index", type=int, help="Inclusive index AFTER split filtering")
|
|
p.add_argument("--limit", type=int, help="First N clips in the selected manifest range")
|
|
p.add_argument("--num-shards", type=int, default=1, help="Disjoint round-robin shards after all sample selection")
|
|
p.add_argument("--shard-index", type=int, default=0, help="Zero-based shard index")
|
|
p.add_argument("--ego4d-annotations", type=Path, nargs="+")
|
|
p.add_argument("--model-id", default="Qwen/Qwen2.5-Omni-3B")
|
|
p.add_argument("--encoder-id", help="Stable checkpoint identity for relocated copies; default: model-id")
|
|
p.add_argument("--device", default="cuda:0")
|
|
p.add_argument("--dtype", choices=("float32", "float16", "bfloat16"), default="bfloat16")
|
|
p.add_argument("--flash-attention-2", action="store_true")
|
|
p.add_argument("--video-fps", type=float, default=2.)
|
|
p.add_argument("--audio-window-seconds", type=float, default=1.)
|
|
p.add_argument("--visual-window-seconds", type=float, default=1.)
|
|
p.add_argument("--box-kind", choices=("face", "head"), default="face")
|
|
p.add_argument(
|
|
"--box-coordinate-space", choices=("camera", "encoded"), default=None,
|
|
help=("Box coordinates before or after video cropping. Default: encoded for EasyCom "
|
|
"(subtract the left banner), camera for uncropped Ego4D. Use camera only for "
|
|
"annotations already translated to the cropped image."),
|
|
)
|
|
p.add_argument("--crop-expansion", type=float, default=1.)
|
|
p.add_argument("--max-pixels", type=int, default=224 * 224)
|
|
p.add_argument("--image-batch-size", type=int, default=8)
|
|
p.add_argument("--decoder-backend", choices=("decord", "torchcodec", "pynvcodec"), default="decord")
|
|
p.add_argument("--decoder-threads", type=int, default=1)
|
|
p.add_argument("--decoder-gpu-id", type=int, default=0, help="Visible-device index, not physical GPU number")
|
|
p.add_argument("--max-steps", type=int, help="Smoke test only; partial sidecars are rejected by training")
|
|
p.add_argument("--preview-only", action="store_true", help="Save box overlays without loading the model")
|
|
p.add_argument("--preview-count", type=int, default=4, help="Maximum images per clip (preview only)")
|
|
p.add_argument("--preview-start", type=float, default=0., help="Start time in seconds on the clip/dense timeline")
|
|
p.add_argument("--preview-end", type=float, help="Exclusive end time; default: entire available prefix")
|
|
p.add_argument("--preview-with-boxes", action="store_true",
|
|
help="Sample only grid frames with valid provided crops; preview selection only, not feature filtering")
|
|
existing = p.add_mutually_exclusive_group()
|
|
existing.add_argument("--overwrite", action="store_true")
|
|
existing.add_argument("--reuse-existing", action="store_true", help="Validate and skip complete matching sidecars")
|
|
p.add_argument("--dry-run", action="store_true", help="Validate references/targets and list work; no model, media decode or writes")
|
|
p.add_argument("--no-progress", action="store_true")
|
|
args = p.parse_args(argv)
|
|
if args.num_shards <= 0 or not 0 <= args.shard_index < args.num_shards:
|
|
p.error("Require --num-shards > 0 and 0 <= --shard-index < --num-shards")
|
|
if args.manifest is not None and args.dataset != "ego4d":
|
|
p.error("--manifest is only supported with --dataset ego4d")
|
|
if args.manifest is None and (args.split != "all" or args.start_index != 1
|
|
or args.end_index is not None or args.limit is not None):
|
|
p.error("--split/--start-index/--end-index/--limit require --manifest")
|
|
if args.start_index < 1 or (args.end_index is not None and args.end_index < args.start_index):
|
|
p.error("Invalid manifest index range")
|
|
if args.limit is not None and args.limit <= 0:
|
|
p.error("--limit must be positive")
|
|
if args.reuse_existing and (args.max_steps is not None or args.preview_only):
|
|
p.error("--reuse-existing is for full extraction, not smoke tests or previews")
|
|
if args.dry_run and args.preview_only:
|
|
p.error("--dry-run and --preview-only are separate checks")
|
|
if args.preview_count <= 0:
|
|
p.error("--preview-count must be positive")
|
|
if not np.isfinite(args.preview_start) or args.preview_start < 0:
|
|
p.error("--preview-start must be finite and nonnegative")
|
|
if args.preview_end is not None and (not np.isfinite(args.preview_end) or args.preview_end <= args.preview_start):
|
|
p.error("--preview-end must be finite and greater than --preview-start")
|
|
if not args.preview_only and (args.preview_count != 4 or args.preview_start != 0.
|
|
or args.preview_end is not None or args.preview_with_boxes):
|
|
p.error("Preview selection options require --preview-only")
|
|
if args.preview_only and args.max_steps is not None:
|
|
if args.preview_start >= args.max_steps / 5 or (args.preview_end is not None
|
|
and args.preview_end > args.max_steps / 5):
|
|
p.error("Preview range exceeds --max-steps; remove --max-steps to inspect later frames")
|
|
if args.box_coordinate_space is None:
|
|
args.box_coordinate_space = "encoded" if args.dataset == "easycom" else "camera"
|
|
for name in ("video_fps", "audio_window_seconds", "visual_window_seconds", "crop_expansion"):
|
|
if not np.isfinite(getattr(args, name)) or getattr(args, name) <= 0:
|
|
p.error(f"--{name.replace('_', '-')} must be finite and positive")
|
|
if args.audio_window_seconds > 30:
|
|
p.error("--audio-window-seconds must not exceed the audio processor's 30s limit")
|
|
if args.max_pixels < 4 * 28 * 28 or args.image_batch_size <= 0:
|
|
p.error("max-pixels must be >=3136; image-batch-size must be positive")
|
|
if args.max_steps is not None and args.max_steps <= 0:
|
|
p.error("--max-steps must be positive")
|
|
protected_roots = [args.feature_root]
|
|
if args.dataset == "ego4d":
|
|
protected_roots += [args.feature_root / split for split in ("train", "test")]
|
|
if args.output_root.resolve() in [root.resolve() for root in protected_roots]:
|
|
p.error("Local outputs must not overwrite original Omni features")
|
|
return args
|
|
|
|
|
|
def main(argv=None) -> None:
|
|
args = parse_args(argv)
|
|
version = None
|
|
if not args.preview_only and not args.dry_run:
|
|
import transformers
|
|
|
|
# Fail before indexing/decoding source media, not after expensive I/O.
|
|
version = validate_local_transformers_version(transformers.__version__)
|
|
names = select_samples(args)
|
|
if not names:
|
|
print(f"Shard {args.shard_index}/{args.num_shards}: no assigned clips")
|
|
return
|
|
print(f"Shard {args.shard_index}/{args.num_shards}: {len(names)} assigned clips", flush=True)
|
|
reference_paths = {name: resolve_feature_path(args.feature_root, name, dataset=args.dataset) for name in names}
|
|
paths = []
|
|
if args.dataset == "ego4d":
|
|
paths = args.ego4d_annotations or [args.data_root / "v1" / "annotations" / f"av_{split}.json"
|
|
for split in ("train", "val")]
|
|
for path in paths:
|
|
if not path.is_file():
|
|
raise FileNotFoundError(path)
|
|
for name in names:
|
|
required = [reference_paths[name], args.target_root / f"{name}.npy"]
|
|
if args.dataset == "ego4d":
|
|
required.append(args.data_root / "v1" / "clips" / f"{name}.mp4")
|
|
for path in required:
|
|
if not path.is_file():
|
|
raise FileNotFoundError(path)
|
|
if ((args.output_root / f"{name}.pt").exists() and not args.overwrite
|
|
and not args.preview_only and not args.reuse_existing):
|
|
raise FileExistsError(f"{name}.pt exists; use --reuse-existing, a new output root or --overwrite")
|
|
pending, contract = [], None
|
|
# Stream through references on CPU; never keep all global states in memory.
|
|
for name in tqdm(names, desc="local preflight", unit="clip", disable=args.no_progress or args.dry_run):
|
|
reference = torch.load(reference_paths[name], map_location="cpu", weights_only=False)
|
|
times, _ = reference_alignment(reference)
|
|
target = np.load(args.target_root / f"{name}.npy", allow_pickle=False, mmap_mode="r")
|
|
if target.shape != (len(times), 3):
|
|
raise ValueError(f"Reference/target mismatch: {name}, {target.shape} vs {len(times)} states")
|
|
if args.reuse_existing and (args.output_root / f"{name}.pt").is_file():
|
|
contract = validate_reusable_sidecar(args, name, reference, version, contract)
|
|
print(f"{name}: reuse ({len(times)} states)", flush=True)
|
|
else:
|
|
pending.append(name)
|
|
if args.dry_run:
|
|
steps = min(len(times), args.max_steps or len(times))
|
|
print(f"{name}: extract ({steps} states, {steps / 5:g}s)", flush=True)
|
|
del reference, target
|
|
print(f"Selected {len(names)} clips; reuse={len(names) - len(pending)}, pending={len(pending)}", flush=True)
|
|
if args.dry_run:
|
|
print("Dry run only: annotation contents, pixel alignment and model/runtime are not checked.")
|
|
return
|
|
if not pending:
|
|
return
|
|
ego_boxes = load_ego4d_boxes(paths, pending) if args.dataset == "ego4d" else {}
|
|
args.output_root.mkdir(parents=True, exist_ok=True)
|
|
encoder = None
|
|
for position, name in enumerate(pending, 1):
|
|
print(f"[{position}/{len(pending)}] {name}", flush=True)
|
|
path = reference_paths[name]
|
|
reference = torch.load(path, map_location="cpu", weights_only=False)
|
|
times, indices = reference_alignment(reference)
|
|
target = np.load(args.target_root / f"{name}.npy", allow_pickle=False, mmap_mode="r")
|
|
if target.shape != (len(times), 3):
|
|
raise ValueError(f"Reference/target mismatch: {name}, {target.shape} vs {len(times)} states")
|
|
full_steps = len(times)
|
|
if args.max_steps is not None:
|
|
times, indices = times[:args.max_steps], indices[:args.max_steps]
|
|
media_end = len(times) / 5
|
|
if args.preview_only:
|
|
_, media_end = preview_time_bounds(args, media_end)
|
|
media, annotations = build_media_and_boxes(args, name, media_end, ego_boxes)
|
|
validate_source_timeline(reference, media)
|
|
if (args.dataset == "easycom" and args.box_coordinate_space == "camera"
|
|
and any(clip.crop is not None and clip.crop[0] for clip in media.clips)):
|
|
print("NOTE: camera boxes receive NO banner-offset correction. For original EasyCom "
|
|
"annotations use --box-coordinate-space encoded; otherwise boxes shift right.", flush=True)
|
|
if args.preview_only:
|
|
preview_boxes(media, annotations, args, name, len(times) / 5)
|
|
if args.dataset == "ego4d":
|
|
audit_path = args.output_root / f"{name}_preview" / "box_frame_audit.json"
|
|
audit_path.write_text(json.dumps(ego_boxes[name]["frame_audit"], indent=2) + "\n")
|
|
ego_boxes.pop(name, None)
|
|
del media, annotations, reference
|
|
continue
|
|
if encoder is None:
|
|
encoder = FrozenQwenLocalEncoder(args)
|
|
config = {**requested_local_config(args),
|
|
"model_config_sha256": encoder.model_config_sha256,
|
|
"transformers": encoder.transformers_version,
|
|
"visual_dim": encoder.visual_dim, "audio_dim": encoder.audio_dim}
|
|
new_contract = {"format": LOCAL_FORMAT, "modalities": "av", "config": config}
|
|
if contract is not None and contract != new_contract:
|
|
raise ValueError("Current encoder differs from reused sidecars; use a new output root")
|
|
contract = new_contract
|
|
print(f"{name}: {len(times)} states; native visual={encoder.visual_dim}D audio={encoder.audio_dim}D", flush=True)
|
|
output = extract_local(media, annotations, encoder, times, args)
|
|
output.update({
|
|
"format": LOCAL_FORMAT, "dataset": args.dataset, "sample_id": name,
|
|
"times": times, "label_indices": indices, "partial": len(times) != full_steps,
|
|
"reference_features": str(path.resolve()), "direct_media": media.metadata(),
|
|
"config": config,
|
|
})
|
|
if args.dataset == "ego4d":
|
|
# Per-clip diagnostics, not encoder config: existing valid sidecars
|
|
# retain the same reuse contract and need not be re-extracted.
|
|
output["box_frame_audit"] = ego_boxes[name]["frame_audit"]
|
|
destination = args.output_root / f"{name}.pt"
|
|
# Atomic replacement: a failed extraction never looks like a ready file.
|
|
with tempfile.NamedTemporaryFile(dir=args.output_root, prefix=f".{name}.", suffix=".tmp", delete=False) as handle:
|
|
temporary = Path(handle.name)
|
|
try:
|
|
torch.save(output, temporary)
|
|
os.replace(temporary, destination)
|
|
finally:
|
|
temporary.unlink(missing_ok=True)
|
|
print(f"Saved {destination}; V-valid={output['visual_valid'].float().mean():.1%}, "
|
|
f"A-valid={output['audio_valid'].float().mean():.1%}, partial={output['partial']}", flush=True)
|
|
ego_boxes.pop(name, None)
|
|
del media, annotations, reference, output
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|