PruneVid for aerial drone videos:

"""Single-clip adapter for Visual-AI/PruneVid. Run from its repository root."""
import argparse
import json
import math
import sys
import time
from pathlib import Path


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument('--video', required=True)
    p.add_argument('--output', default='drone_result')
    p.add_argument('--start', type=float, default=0)
    p.add_argument('--seconds', type=float, default=20)
    p.add_argument('--frames', type=int, default=16)
    p.add_argument('--prepare-only', action='store_true')
    p.add_argument('--model-dir', default='MODELS/pllava-7b')
    p.add_argument('--weight-dir')
    p.add_argument('--lora-alpha', type=int, default=14)
    p.add_argument('--question', default='Describe the visible terrain, structures, roads and moving objects. Distinguish camera motion from object motion. State uncertainty; do not invent GPS coordinates or metric distances.')
    p.add_argument('--alpha', type=float, default=0.4)
    p.add_argument('--tau', type=float, default=0.8)
    p.add_argument('--cluster-ratio', type=float, default=0.5)
    p.add_argument('--temporal-segment-ratio', type=float, default=0.25)
    p.add_argument('--selected-layer', type=int, default=10)
    p.add_argument('--max-new-tokens', type=int, default=160)
    a = p.parse_args()
    if a.frames < 4 or a.seconds <= 0 or a.start < 0:
        p.error('Require frames >= 4, seconds > 0 and start >= 0.')
    for name in ('alpha', 'cluster_ratio', 'temporal_segment_ratio'):
        if not 0 < getattr(a, name) <= 1:
            p.error(name + ' must be in (0, 1].')
    if not -1 <= a.tau <= 1:
        p.error('tau must be in [-1, 1].')
    import cv2
    import numpy as np
    from PIL import Image, ImageDraw
    video = Path(a.video).resolve()
    if not video.is_file():
        raise FileNotFoundError(video)
    target = Path(a.output)
    target.mkdir(parents=True, exist_ok=True)
    cap = cv2.VideoCapture(str(video))
    try:
        if not cap.isOpened():
            raise RuntimeError('Cannot open video; check path and codec.')
        fps = cap.get(cv2.CAP_PROP_FPS)
        total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        if not math.isfinite(fps) or fps <= 0 or total <= 0:
            raise RuntimeError('Invalid FPS/frame count; transcode video before sampling.')
        first = int(a.start * fps)
        last = min(total, int((a.start + a.seconds) * fps)) - 1
        if last - first + 1 < a.frames:
            raise RuntimeError('Requested clip has fewer source frames than --frames.')
        indices = np.linspace(first, last, a.frames).round().astype(int)
        images, samples = [], []
        for i, idx in enumerate(indices):
            cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
            ok, bgr = cap.read()
            if not ok:
                raise RuntimeError(f'Cannot decode source frame {idx}.')
            image = Image.fromarray(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
            images.append(image)
            image.save(target / f'frame_{i:03d}.jpg')
            samples.append({'frame_index': int(idx), 'nominal_seconds': float(idx / fps)})
    finally:
        cap.release()
    sheet = Image.new('RGB', (4 * 320, math.ceil(a.frames / 4) * 205), 'white')
    draw = ImageDraw.Draw(sheet)
    for i, image in enumerate(images):
        thumb = image.copy()
        thumb.thumbnail((320, 180))
        x, y = (i % 4) * 320, (i // 4) * 205
        sheet.paste(thumb, (x, y))
        draw.text((x + 4, y + 182), f'{samples[i]["nominal_seconds"]:.2f}s', fill='black')
    sheet.save(target / 'contact_sheet.jpg')
    metadata = {'arguments': vars(a), 'video': str(video), 'fps': fps,
                'timestamps_are_nominal': True, 'samples': samples}
    (target / 'sampling.json').write_text(json.dumps(metadata, indent=2), encoding='utf-8')
    if a.prepare_only:
        print('Prepared sampled frames and contact sheet; no model loaded.')
        return
    import torch
    if not torch.cuda.is_available():
        raise RuntimeError('Official-model path requires working NVIDIA CUDA in this adapter. Use --prepare-only for CPU preprocessing.')
    if not torch.cuda.is_bf16_supported():
        raise RuntimeError('Official loader uses bfloat16; this CUDA device does not report BF16 support.')
    root = Path.cwd()
    if not (root / 'tasks/eval/model_utils.py').is_file():
        raise RuntimeError('Run from the cloned PruneVid root, with this script copied there.')
    sys.path.insert(0, str(root))
    from tasks.eval.model_utils import load_pllava
    from tasks.eval.eval_utils import conv_templates
    model_dir = str(Path(a.model_dir).resolve())
    weight_dir = str(Path(a.weight_dir or a.model_dir).resolve())
    if 'llava' not in model_dir:
        raise RuntimeError('Official loader selects PLLaVA by lowercase "llava" in model path; name the directory pllava-7b.')
    for directory in (model_dir, weight_dir):
        if not (Path(directory) / 'config.json').is_file():
            raise FileNotFoundError(f'Missing config.json in {directory}')
    torch.manual_seed(0)
    model, processor = load_pllava(
        model_dir, num_frames=a.frames, use_lora=True, weight_dir=weight_dir,
        lora_alpha=a.lora_alpha, pooling_shape=(a.frames, 12, 12),
        selected_layer=a.selected_layer, alpha=a.alpha, tau=a.tau,
        cluster_ratio=a.cluster_ratio, temporal_segment_ratio=a.temporal_segment_ratio)
    model = model.to('cuda').eval()
    conv = conv_templates['plain'].copy()
    conv.user_query(a.question, is_mm=True)
    inputs = processor(text=conv.get_prompt(), images=images, return_tensors='pt').to(model.device)
    if 'pixel_values' in inputs and inputs['pixel_values'] is not None:
        inputs['pixel_values'] = inputs['pixel_values'].to(dtype=next(model.parameters()).dtype)
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    started = time.perf_counter()
    with torch.inference_mode():
        tokens = model.generate(**inputs, media_type='video', do_sample=False,
                                max_new_tokens=a.max_new_tokens, use_cache=True)
    torch.cuda.synchronize()
    decoded = processor.batch_decode(tokens, skip_special_tokens=True,
                                     clean_up_tokenization_spaces=False)[0]
    answer = decoded.split(conv.roles[-1])[-1].removesuffix('</s>').strip()
    result = {**metadata, 'answer': answer, 'raw_decoded_text': decoded,
              'generation_seconds': time.perf_counter() - started,
              'peak_allocated_gpu_gib': torch.cuda.max_memory_allocated() / 2**30,
              'gpu': torch.cuda.get_device_name(0), 'torch_version': torch.__version__,
              'note': 'No measured retained-token count; no baseline or accuracy claim.'}
    (target / 'result.json').write_text(json.dumps(result, indent=2), encoding='utf-8')
    print(answer)


if __name__ == '__main__':
    main()