Sample frames from a video effectively.
# type: ignore[attr-defined]
"""Video sampler allows you to efficiently sample video frames"""
from importlib import metadata as importlib_metadata


def get_version() -> str:
    try:
        return importlib_metadata.version(__name__)
    except importlib_metadata.PackageNotFoundError:  # pragma: no cover
        return "unknown"


version: str = get_version()

video_sampler/__init__.py
LemurPwned/video-sampler
# type: ignore[attr-defined]

import glob
import os
from enum import Enum
from typing import Annotated

import typer
from tqdm import tqdm

from video_sampler import version
from video_sampler.buffer import SamplerConfig, check_args_validity
from video_sampler.logging import Color, console
from video_sampler.sampler import Worker

app = typer.Typer(
    name="video-sampler",
    help="Video sampler allows you to efficiently sample video frames",
    add_completion=True,
)


class BufferType(str, Enum):
    entropy = "entropy"
    gzip = "gzip"
    hash = "hash"
    passthrough = "passthrough"


def version_callback(print_version: bool = True) -> None:
    """Print the version of the package."""
    if print_version:
        console.print(f"[yellow]video-sampler[/] version: [bold blue]{version}[/]")
        raise typer.Exit()


def delegate_workers(video_path: str, output_path: str, cfg: SamplerConfig):
    msg = "Detected input as a file"
    if not os.path.isfile(video_path):
        if "*" not in video_path:
            videos = glob.glob(os.path.join(video_path, "*"))
        else:
            videos = glob.glob(video_path)
        msg = f"Detected input as a folder with {len(videos)} files"
    else:
        videos = [video_path]
    console.print(msg, style=f"bold {Color.cyan.value}")

    worker = Worker(
        cfg=cfg,
    )
    for video in tqdm(videos, desc="Processing videos..."):
        video_subpath = os.path.join(output_path, os.path.basename(video))
        worker.launch(
            video_path=video,
            output_path=video_subpath,
        )


@app.command(name="hash")
def main(
    video_path: str = typer.Argument(
        ..., help="Path to the video file or a glob pattern."
    ),
    output_path: str = typer.Argument(..., help="Path to the output folder."),
    min_frame_interval_sec: float = typer.Option(
        1.0, help="Minimum frame interval in seconds."
    ),
    stats: bool = typer.Option(True, help="Print stats."),
    keyframes_only: bool = typer.Option(True, help="Only sample keyframes."),
    buffer_size: int = typer.Option(10, help="Size of the buffer."),
    hash_size: int = typer.Option(4, help="Size of the hash."),
    queue_wait: float = typer.Option(0.1, help="Time to wait for the queue."),
    debug: bool = typer.Option(False, help="Enable debug mode."),
) -> None:
    """Default buffer is the perceptual hash buffer"""

    cfg = SamplerConfig(
        min_frame_interval_sec=min_frame_interval_sec,
        keyframes_only=keyframes_only,
        queue_wait=queue_wait,
        print_stats=stats,
        debug=debug,
        buffer_config={
            "type": "hash",
            "size": buffer_size,
            "debug": debug,
            "hash_size": hash_size,
        },
    )
    console.print(cfg, style=f"bold {Color.yellow.value}")
    delegate_workers(video_path=video_path, output_path=output_path, cfg=cfg)


@app.command(name="buffer")
def buffer(
    buffer_type: Annotated[BufferType, typer.Argument(case_sensitive=False)],
    video_path: str = typer.Argument(
        ..., help="Path to the video file or a glob pattern."
    ),
    output_path: str = typer.Argument(..., help="Path to the output folder."),
    min_frame_interval_sec: float = typer.Option(
        1.0, help="Minimum frame interval in seconds."
    ),
    stats: bool = typer.Option(True, help="Print stats."),
    keyframes_only: bool = typer.Option(True, help="Only sample keyframes."),
    buffer_size: int = typer.Option(10, help="Size of the buffer."),
    hash_size: int = typer.Option(4, help="Size of the hash."),
    expiry: int = typer.Option(4, help="Expiry time for the buffer."),
    queue_wait: float = typer.Option(0.1, help="Time to wait for the queue."),
    debug: bool = typer.Option(False, help="Enable debug mode."),
):
    """Buffer type can be one of entropy, gzip, hash, passthrough"""
    cfg = SamplerConfig(
        min_frame_interval_sec=min_frame_interval_sec,
        keyframes_only=keyframes_only,
        queue_wait=queue_wait,
        print_stats=stats,
        debug=debug,
        buffer_config={
            "type": buffer_type,
            "size": buffer_size,
            "debug": debug,
            "hash_size": hash_size,
            "expiry": expiry,
        },
    )
    # create a test buffer
    try:
        check_args_validity(cfg)
    except AssertionError as e:
        console.print(
            "Error while creating buffer",
            f"\n\t{e}",
            style=f"bold {Color.red.value}",
        )
        raise typer.Exit(code=-1)

    console.print(cfg, style=f"bold {Color.yellow.value}")
    delegate_workers(video_path=video_path, output_path=output_path, cfg=cfg)


def main_loop():
    app()


if __name__ == "__main__":
    main_loop()

video_sampler/__main__.py
LemurPwned/video-sampler
import gzip
import heapq
from abc import ABC, abstractmethod
from collections import OrderedDict
from collections.abc import Iterable
from dataclasses import asdict, dataclass, field
from typing import Any

from imagehash import average_hash, phash
from PIL import Image

from .logging import Color, console


@dataclass
class SamplerConfig:
    min_frame_interval_sec: float = 1
    keyframes_only: bool = True
    queue_wait: float = 0.1
    debug: bool = False
    print_stats: bool = False
    buffer_config: dict[str, Any] = field(
        default_factory=lambda: {
            "type": "entropy",
            "size": 15,
            "debug": True,
        }
    )

    def __str__(self) -> str:
        return str(asdict(self))


class FrameBuffer(ABC):
    @abstractmethod
    def add(self, item: Image.Image, metadata: dict[str, Any]) -> None | tuple:
        pass

    @abstractmethod
    def final_flush(self) -> Iterable[tuple[Image.Image | None, dict]]:
        """Flush the buffer and return the remaining items"""
        pass

    @abstractmethod
    def get_buffer_state(self) -> list[str]:
        """Return the current state of the buffer"""
        pass

    @abstractmethod
    def clear(self):
        """Clear the buffer"""
        pass


class PassThroughBuffer(FrameBuffer):
    def __init__(self) -> None:
        ...

    def get_buffer_state(self) -> list[str]:
        return []

    def add(self, item: Image.Image, metadata: dict[str, Any]):
        return (item, metadata)

    def final_flush(self) -> Iterable[tuple[Image.Image | None, dict]]:
        yield None

    def clear(self):
        pass


class HashBuffer(FrameBuffer):
    def __init__(self, size: int, debug_flag: bool = False, hash_size: int = 4) -> None:
        self.ordered_buffer = OrderedDict()
        self.max_size = size
        self.debug_flag = debug_flag
        self.hash_size = hash_size

    def get_buffer_state(self) -> list[str]:
        return list(self.ordered_buffer.keys())

    def add(self, item: Image.Image, metadata: dict[str, Any]):
        hash_ = str(phash(item, hash_size=self.hash_size))
        if not self.__check_duplicate(hash_):
            return self.__add(item, hash_, metadata)
        return None

    def __add(self, item: Image.Image, hash_: str, metadata: dict):
        self.ordered_buffer[hash_] = (item, metadata)
        if len(self.ordered_buffer) >= self.max_size:
            return self.ordered_buffer.popitem(last=False)[1]
        return None

    def __check_duplicate(self, hash_: str) -> bool:
        if hash_ in self.ordered_buffer:
            # renew the hash validity
            if self.debug_flag:
                console.print(
                    f"Renewing {hash_}",
                    style=f"bold {Color.red.value}",
                )
            self.ordered_buffer.move_to_end(hash_)
            return True
        return False

    def final_flush(self) -> Iterable[tuple[Image.Image | None, dict]]:
        yield from self.ordered_buffer.values()

    def clear(self):
        self.ordered_buffer.clear()


class SlidingTopKBuffer(FrameBuffer):
    def __init__(
        self, size: int, debug_flag: bool = False, expiry: int = 30, hash_size: int = 8
    ) -> None:
        # it's a min heap with fixed size
        self.sliding_buffer = []
        self.max_size = size
        self.debug_flag = debug_flag
        self.expiry_count = expiry
        self.hash_size = hash_size
        assert (
            self.expiry_count > self.max_size
        ), "expiry count must be greater than max size"
        console.print(
            f"Creating sliding buffer of size {self.max_size} and expiry {expiry}",
            style=f"bold {Color.red.value}",
        )

    def get_buffer_state(self) -> list[str]:
        return [item[:3] for item in self.sliding_buffer]

    def add(self, item: Image.Image, metadata: dict[str, Any]):
        assert "index" in metadata, "metadata must have index key for sliding buffer"
        average_hash_ = str(average_hash(item, hash_size=self.hash_size))
        to_return = None
        if not self.__check_duplicate(average_hash_):
            heapq.heappush(
                self.sliding_buffer,
                [metadata["index"], 0, average_hash_, item, metadata],
            )
            if len(self.sliding_buffer) >= self.max_size:
                to_return = heapq.heappop(self.sliding_buffer)[-2:]
        # update the expiry count
        expired_indx = -1
        for i in range(len(self.sliding_buffer)):
            self.sliding_buffer[i][1] += 1
            if self.sliding_buffer[i][1] >= self.expiry_count:
                expired_indx = i
        # at any point only one item can be expired
        if expired_indx != -1:
            self.sliding_buffer.pop(expired_indx)  # just drop
        return to_return

    def __check_duplicate(self, hash_: str) -> bool:
        for item in self.sliding_buffer:
            if item[2] == hash_:
                # renew the hash validity
                if self.debug_flag:
                    console.print(
                        f"Renewing {hash_}",
                        style=f"bold {Color.red.value}",
                    )
                item[1] = 0
                return True
        return False

    def final_flush(self) -> Iterable[tuple[Image.Image | None, dict]]:
        if len(self.sliding_buffer):
            yield heapq.heappop(self.sliding_buffer)[-2:]
        yield None, {}

    def clear(self):
        self.sliding_buffer.clear()


class GzipBuffer(FrameBuffer):
    """Measure compression size as a function of the image usability"""

    def __init__(
        self, size: int, expiry: int, debug_flag: bool = False, hash_size: int = 8
    ) -> None:
        self.sliding_top_k_buffer = SlidingTopKBuffer(
            size=size, expiry=expiry, debug_flag=debug_flag, hash_size=hash_size
        )

    def get_buffer_state(self) -> list[str]:
        return self.sliding_top_k_buffer.get_buffer_state()

    def add(self, item: Image.Image, metadata: dict[str, Any]):
        compressed_l = len(gzip.compress(item.tobytes()))
        return self.sliding_top_k_buffer.add(item, {**metadata, "index": -compressed_l})

    def final_flush(self) -> Iterable[tuple[Image.Image | None, dict]]:
        return self.sliding_top_k_buffer.final_flush()

    def clear(self):
        self.sliding_top_k_buffer.clear()


class EntropyByffer(FrameBuffer):
    """Measure image entropy as a function of the image usability"""

    def __init__(
        self, size: int, expiry: int, debug_flag: bool = False, hash_size: int = 8
    ) -> None:
        self.sliding_top_k_buffer = SlidingTopKBuffer(
            size=size, expiry=expiry, debug_flag=debug_flag, hash_size=hash_size
        )

    def get_buffer_state(self) -> list[str]:
        return self.sliding_top_k_buffer.get_buffer_state()

    def add(self, item: Image.Image, metadata: dict[str, Any]):
        entropy = item.entropy()
        return self.sliding_top_k_buffer.add(item, {**metadata, "index": -entropy})

    def final_flush(self) -> Iterable[tuple[Image.Image | None, dict]]:
        return self.sliding_top_k_buffer.final_flush()

    def clear(self):
        self.sliding_top_k_buffer.clear()


def check_args_validity(cfg: SamplerConfig):
    assert (
        cfg.min_frame_interval_sec > 0
    ), "min_frame_interval_sec must be greater than 0"
    assert cfg.buffer_config["size"] > 0, "buffer size must be greater than 0"
    arg_check = {
        "hash": ("hash_size", "size"),
        "gzip": ("hash_size", "size", "expiry"),
        "entropy": ("hash_size", "size", "expiry"),
    }
    for arg in arg_check[cfg.buffer_config["type"]]:
        assert arg in cfg.buffer_config, f"{arg} must be present in buffer config"
        assert (
            cfg.buffer_config[arg] is not None and cfg.buffer_config[arg] > 0
        ), f"{arg} must be greater than 0 and must not be None"


def create_buffer(buffer_config: dict[str, Any]):
    """Create a buffer based on the config"""
    console.print(
        f"Creating buffer of type {buffer_config['type']}",
        style=f"bold {Color.red.value}",
    )
    if buffer_config["type"] == "hash":
        return HashBuffer(
            size=buffer_config["size"],
            debug_flag=buffer_config["debug"],
            hash_size=buffer_config["hash_size"],
        )
    elif buffer_config["type"] == "sliding_top_k":
        return SlidingTopKBuffer(
            size=buffer_config["size"],
            debug_flag=buffer_config["debug"],
            expiry=buffer_config["expiry"],
        )
    elif buffer_config["type"] == "passthrough":
        return PassThroughBuffer()
    elif buffer_config["type"] == "gzip":
        return GzipBuffer(
            size=buffer_config["size"],
            debug_flag=buffer_config["debug"],
            hash_size=buffer_config["hash_size"],
            expiry=buffer_config["expiry"],
        )
    elif buffer_config["type"] == "entropy":
        return EntropyByffer(
            size=buffer_config["size"],
            debug_flag=buffer_config["debug"],
            hash_size=buffer_config["hash_size"],
            expiry=buffer_config["expiry"],
        )
    else:
        raise ValueError(f"Unknown buffer type {buffer_config['type']}")

video_sampler/buffer.py
LemurPwned/video-sampler
from enum import Enum

from rich.console import Console

console = Console()


class Color(str, Enum):
    white = "white"
    red = "red"
    cyan = "cyan"
    magenta = "magenta"
    yellow = "yellow"
    green = "green"

video_sampler/logging.py
LemurPwned/video-sampler
import os
import time
from collections import Counter
from collections.abc import Iterable
from queue import Queue
from threading import Thread

import av
from PIL import Image

from .buffer import SamplerConfig, create_buffer
from .logging import Color, console


class VideoSampler:
    def __init__(self, cfg: SamplerConfig) -> None:
        self.cfg = cfg
        self.frame_buffer = create_buffer(self.cfg.buffer_config)
        self.stats = Counter()

    def sample(self, video_path: str) -> Iterable[tuple[Image.Image | None, dict]]:
        """Generate sample frames from a video"""
        self.stats.clear()
        self.frame_buffer.clear()
        with av.open(video_path) as container:
            stream = container.streams.video[0]
            if self.cfg.keyframes_only:
                stream.codec_context.skip_frame = "NONKEY"
            prev_time = -10
            for frame_indx, frame in enumerate(container.decode(stream)):
                # skip frames if keyframes_only is True
                time_diff = frame.time - prev_time
                self.stats["total"] += 1
                if time_diff < self.cfg.min_frame_interval_sec:
                    continue
                prev_time = frame.time

                frame_pil: Image = frame.to_image()
                if self.cfg.debug:
                    buf = self.frame_buffer.get_buffer_state()
                    console.print(
                        f"Frame {frame_indx}\ttime: {frame.time}",
                        f"\t Buffer ({len(buf)}): {buf}",
                        style=f"bold {Color.green.value}",
                    )
                frame_meta = {"frame_time": frame.time, "frame_indx": frame_indx}
                res = self.frame_buffer.add(
                    frame_pil,
                    metadata=frame_meta,
                )
                self.stats["decoded"] += 1
                if res:
                    self.stats["produced"] += 1
                    yield res

        # flush buffer
        for res in self.frame_buffer.final_flush():
            if res:
                self.stats["produced"] += 1
                yield res

        yield None, {"end": True}

    def write_queue(self, video_path: str, q: Queue):
        try:
            for item in self.sample(video_path=video_path):
                q.put(item)
        except (av.IsADirectoryError, av.InvalidDataError) as e:
            console.print(
                f"Error while processing {video_path}",
                f"\n\t{e}",
                style=f"bold {Color.red.value}",
            )
            q.put(None, {"end": True})


class Worker:
    def __init__(self, cfg: SamplerConfig, devnull: bool = False) -> None:
        self.cfg = cfg
        self.processor = VideoSampler(cfg=cfg)
        self.q = Queue()
        self.devnull = devnull

    def launch(self, video_path: str, output_path: str) -> None:
        os.makedirs(output_path, exist_ok=True)
        proc_thread = Thread(
            target=self.processor.write_queue, args=(video_path, self.q)
        )
        proc_thread.start()
        self.queue_reader(output_path, read_interval=self.cfg.queue_wait)
        proc_thread.join()
        if self.cfg.print_stats:
            console.print(
                f"Stats for: {os.path.basename(video_path)}",
                f"\n\tTotal frames: {self.processor.stats['total']}",
                f"\n\tDecoded frames: {self.processor.stats['decoded']}",
                f"\n\tProduced frames: {self.processor.stats['produced']}",
                style=f"bold {Color.magenta.value}",
            )

    def queue_reader(self, output_path, read_interval=0.1) -> None:
        while True:
            if not self.q.empty():
                item = self.q.get()
                frame, metadata = item
                if frame is not None:
                    if not self.devnull and isinstance(frame, Image.Image):
                        frame.save(
                            os.path.join(output_path, f"{metadata['frame_time']}.jpg")
                        )
                if metadata.get("end", False):
                    break
            time.sleep(read_interval)

video_sampler/sampler.py
LemurPwned/video-sampler
from collections import defaultdict
from itertools import islice
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from sklearn.cluster import KMeans
from sklearn.manifold import TSNE
from tqdm import tqdm
from transformers import AutoFeatureExtractor, ResNetModel


def batched(iterable, n):
    "Batch data into tuples of length n. The last batch may be shorter."
    """from https://docs.python.org/3/library/itertools.html#itertools-recipes"""
    if n < 1:
        raise ValueError("n must be at least one")
    it = iter(iterable)
    while batch := tuple(islice(it, n)):
        yield batch


def build_feature_model(model_str: str):
    extractor = AutoFeatureExtractor.from_pretrained(model_str)
    model = ResNetModel.from_pretrained(model_str)
    return model, extractor


def extract_features(
    model_str: str, image_folder: Path, mkey="pixel_values", batch_size: int = 8
):
    out_features = defaultdict(list)
    model, extractor = build_feature_model(model_str)
    with torch.no_grad():
        all_files = list(image_folder.iterdir())
        for batch in tqdm(
            batched(all_files, batch_size), total=len(all_files) // batch_size
        ):
            # load images
            batch_imgs = [Image.open(img_path).convert("RGB") for img_path in batch]
            # extract features
            batch_imgs = extractor(batch_imgs, return_tensors="pt")[mkey]
            batch_features = model(batch_imgs).pooler_output.squeeze()
            if len(batch) == 1:
                batch_features = batch_features.expand(1, -1)
            batch_features = torch.functional.F.normalize(batch_features, p=2, dim=1)
            out_features["embeds"].extend(batch_features)
            out_features["paths"].extend([img_path.name for img_path in batch])
    return out_features


def cluster_features(
    features,
    max_clusters=50,
):
    proj = TSNE(n_components=2, perplexity=35, metric="cosine")
    Xorg = np.asarray(features["embeds"])
    X = proj.fit_transform(Xorg)

    # take about 10% of the frame as the number of clusters
    n_clusters = min(int(0.1 * len(features["embeds"])), max_clusters)
    cluster_model = KMeans(n_clusters=n_clusters, random_state=0).fit(Xorg)
    return X, cluster_model.labels_

video_sampler/visualisation/clustering.py
LemurPwned/video-sampler