icassp2027-experiment/extract_turn_local_features.py

"""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 normalize_sample_name
from local_turn_features import (
    LOCAL_FORMAT, causal_frame_ticks, crop_box, easycom_boxes,
    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 = {}
    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 names:
                    if uid in result:
                        raise ValueError(f"Duplicate Ego4D annotation UID: {uid}")
                    result[uid] = ego4d_clip_boxes(clip)
        del payload
    missing = set(names) - result.keys()
    if missing:
        raise ValueError(f"No provided tracks for: {sorted(missing)}")
    return result


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:
        annotations = [ego_boxes[name]]
        if annotations[0] and max(annotations[0]) >= media.clips[0].num_frames:
            raise ValueError(f"Ego4D annotation frame exceeds source clip: {name}")
    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_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)
    ticks = list(causal_frame_ticks(end, args.video_fps))
    audit = []
    for index in np.linspace(0, len(ticks) - 1, min(4, len(ticks)), dtype=int):
        time = ticks[index] / 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")
        filename = f"frame_{ticks[index]: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}")
    (folder / "boxes.json").write_text(json.dumps(audit, indent=2) + "\n")
    print(f"Box previews: {folder} (inspect before extracting all sessions)")


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 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)
    p.add_argument("--target-root", required=True, type=Path)
    p.add_argument("--output-root", required=True, type=Path)
    p.add_argument("--samples", required=True, nargs="+", help="EasyCom session numbers or Ego4D clip UIDs")
    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("--overwrite", action="store_true")
    p.add_argument("--no-progress", action="store_true")
    args = p.parse_args(argv)
    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")
    if args.output_root.resolve() == args.feature_root.resolve():
        p.error("Local outputs must not overwrite original Omni features")
    return args


def main() -> None:
    args = parse_args()
    if not args.preview_only:
        import transformers

        # Fail before indexing/decoding source media, not after expensive I/O.
        validate_local_transformers_version(transformers.__version__)
    names = [normalize_sample_name(value, args.dataset) for value in args.samples]
    if len(set(names)) != len(names):
        raise ValueError("Duplicate samples")
    for name in names:
        for path in (args.feature_root / f"{name}.pt", args.target_root / f"{name}.npy"):
            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:
            raise FileExistsError(f"{name}.pt exists; use a new output root or --overwrite")
    ego_boxes = {}
    if args.dataset == "ego4d":
        paths = args.ego4d_annotations or [args.data_root / "v1" / "annotations" / f"av_{split}.json"
                                          for split in ("train", "val")]
        ego_boxes = load_ego4d_boxes(paths, names)
    args.output_root.mkdir(parents=True, exist_ok=True)
    encoder = None
    for name in names:
        path = args.feature_root / f"{name}.pt"
        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, annotations = build_media_and_boxes(args, name, len(times) / 5, 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)
            continue
        if encoder is None:
            encoder = FrozenQwenLocalEncoder(args)
        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": {
                "encoder_id": args.encoder_id or args.model_id,
                "model_config_sha256": encoder.model_config_sha256,
                "transformers": encoder.transformers_version, "dtype": args.dtype,
                "attention": "flash_attention_2" if args.flash_attention_2 else "default",
                "visual_dim": encoder.visual_dim, "audio_dim": encoder.audio_dim,
                "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",
            },
        })
        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)


if __name__ == "__main__":
    main()