File size: 2,130 Bytes
938e515 |
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 |
# Copyright (c) Facebook, Inc. and its affiliates.
import random
from collections.abc import Callable
from enum import Enum
from typing import Callable as TCallable
from typing import List
FrameTsList = List[int]
FrameSelector = TCallable[[FrameTsList], FrameTsList]
class FrameSelectionStrategy(Enum):
"""
Frame selection strategy used with videos:
- "random_k": select k random frames
- "first_k": select k first frames
- "last_k": select k last frames
- "all": select all frames
"""
# fmt: off
RANDOM_K = "random_k"
FIRST_K = "first_k"
LAST_K = "last_k"
ALL = "all"
# fmt: on
class RandomKFramesSelector(Callable): # pyre-ignore[39]
"""
Selector that retains at most `k` random frames
"""
def __init__(self, k: int):
self.k = k
def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
"""
Select `k` random frames
Args:
frames_tss (List[int]): timestamps of input frames
Returns:
List[int]: timestamps of selected frames
"""
return random.sample(frame_tss, min(self.k, len(frame_tss)))
class FirstKFramesSelector(Callable): # pyre-ignore[39]
"""
Selector that retains at most `k` first frames
"""
def __init__(self, k: int):
self.k = k
def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
"""
Select `k` first frames
Args:
frames_tss (List[int]): timestamps of input frames
Returns:
List[int]: timestamps of selected frames
"""
return frame_tss[: self.k]
class LastKFramesSelector(Callable): # pyre-ignore[39]
"""
Selector that retains at most `k` last frames from video data
"""
def __init__(self, k: int):
self.k = k
def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
"""
Select `k` last frames
Args:
frames_tss (List[int]): timestamps of input frames
Returns:
List[int]: timestamps of selected frames
"""
return frame_tss[-self.k :]
|