playback/engine.py

import threading
import time
from collections import deque

import torch

from mediatok.pipeline.decoder import DecoderPipeline


class FrameQueue:
    def __init__(self, maxsize: int = 4):
        self._q: deque[torch.Tensor] = deque(maxlen=maxsize)
        self._lock = threading.Lock()
        self._cv = threading.Condition(self._lock)
        self._closed = False

    def put(self, frame: torch.Tensor):
        with self._lock:
            if self._closed:
                return
            if len(self._q) == self._q.maxlen:
                self._q.popleft()
            self._q.append(frame)
            self._cv.notify()

    def get(self, timeout: float = None) -> torch.Tensor:
        with self._lock:
            while not self._q and not self._closed:
                if not self._cv.wait(timeout=timeout):
                    raise TimeoutError("frame queue timeout")
            if self._closed and not self._q:
                raise StopIteration()
            return self._q.popleft()

    def close(self):
        with self._lock:
            self._closed = True
            self._cv.notify_all()

    @property
    def qsize(self) -> int:
        with self._lock:
            return len(self._q)


class Player:
    def __init__(self, pipeline: DecoderPipeline, mode: str = "realtime",
                 layer_mask: int = 0b111111, queue_depth: int = 4):
        self.pipeline = pipeline
        self.mode = mode
        self.layer_mask = layer_mask
        self.queue = FrameQueue(maxsize=queue_depth)
        self._decode_thread: threading.Thread = None
        self._running = False

    def start(self):
        self._running = True
        self._decode_thread = threading.Thread(target=self._decode_loop, daemon=True)
        self._decode_thread.start()

    def stop(self):
        self._running = False
        self.queue.close()

    def _decode_loop(self):
        if self.mode == "predecode":
            all_frames = self.pipeline.decode_all(layer_mask=self.layer_mask)
            for f in all_frames:
                self.queue.put(f.cpu())
                if not self._running:
                    break
        elif self.mode == "progressive":
            for i in range(self.pipeline.reader.num_chunks):
                frame = self.pipeline.decode_chunk(i, layer_mask=self.layer_mask)
                self.queue.put(frame.cpu())
                if not self._running:
                    break
        else:
            fps = self.pipeline.reader.header.fps
            frame_time = 1.0 / fps if fps > 0 else 0.033
            for i in range(self.pipeline.reader.num_chunks):
                t0 = time.perf_counter()
                frame = self.pipeline.decode_chunk(i, layer_mask=self.layer_mask)
                self.queue.put(frame.cpu())
                elapsed = time.perf_counter() - t0
                sleep_time = frame_time - elapsed
                if sleep_time > 0:
                    time.sleep(sleep_time)
                if not self._running:
                    break
        self.queue.close()

    def play(self):
        self.start()
        try:
            while True:
                frame = self.queue.get()
                self._display_frame(frame)
        except (StopIteration, TimeoutError):
            pass
        finally:
            self.stop()

    def _display_frame(self, frame: torch.Tensor):
        pass