mirror of https://github.com/RE-OWOD/RE-OWOD
88 lines
2.0 KiB
Python
88 lines
2.0 KiB
Python
|
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
|
||
|
|
|
||
|
|
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):
|
||
|
|
"""
|
||
|
|
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):
|
||
|
|
"""
|
||
|
|
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):
|
||
|
|
"""
|
||
|
|
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 :]
|