Spaces:
Paused
Paused
File size: 3,382 Bytes
92076a7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | 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
|