playback/display.py

import time

import numpy as np
import torch
from PIL import Image


def tensor_to_image(frame: torch.Tensor) -> Image.Image:
    """Convert [3, H, W] or [1, 3, H, W] float tensor to PIL Image."""
    f = frame.detach().cpu().float()
    if f.dim() == 4:
        f = f.squeeze(0)
    arr = f.permute(1, 2, 0).numpy()
    arr = ((arr - arr.min()) / (arr.max() - arr.min() + 1e-8) * 255).astype(np.uint8)
    return Image.fromarray(arr)


class TkPlayer:
    """Simple tkinter video player.

    Decodes chunks through the pipeline and displays frames in a
    window at the source frame rate.  Close the window to stop.
    """
    def __init__(self, pipeline, layer_mask: int = 0b111111):
        self.pipeline = pipeline
        self.layer_mask = layer_mask

    def play(self):
        reader = self.pipeline.reader
        fps = reader.header.fps or 30
        frame_time = 1.0 / fps

        frames = []
        for chunk in self.pipeline.decode_all(layer_mask=self.layer_mask):
            for t in range(chunk.shape[2]):
                frames.append(chunk[:, :, t])
        if not frames:
            raise RuntimeError("no frames decoded")

        import tkinter as tk
        from PIL import ImageTk

        root = tk.Tk()
        root.title(f"MediaTok — {reader.path}")
        label = tk.Label(root)
        label.pack()
        first = tensor_to_image(frames[0])
        photo = ImageTk.PhotoImage(first)
        label.configure(image=photo)
        root.geometry(f"{first.width}x{first.height}")

        state = {"index": 1, "photo": photo, "last": time.perf_counter()}

        def advance():
            if state["index"] >= len(frames):
                root.destroy()
                return
            img = tensor_to_image(frames[state["index"]])
            state["photo"] = ImageTk.PhotoImage(img)
            label.configure(image=state["photo"])
            state["index"] += 1
            elapsed = time.perf_counter() - state["last"]
            state["last"] = time.perf_counter()
            delay = max(1, int((frame_time - elapsed) * 1000))
            root.after(delay, advance)

        root.after(0, advance)
        root.mainloop()