|
|
|
import os.path as osp
|
|
from collections import OrderedDict
|
|
|
|
import cv2
|
|
from cv2 import (CAP_PROP_FOURCC, CAP_PROP_FPS, CAP_PROP_FRAME_COUNT,
|
|
CAP_PROP_FRAME_HEIGHT, CAP_PROP_FRAME_WIDTH,
|
|
CAP_PROP_POS_FRAMES, VideoWriter_fourcc)
|
|
|
|
from annotator.uniformer.mmcv.utils import (check_file_exist, mkdir_or_exist, scandir,
|
|
track_progress)
|
|
|
|
|
|
class Cache:
|
|
|
|
def __init__(self, capacity):
|
|
self._cache = OrderedDict()
|
|
self._capacity = int(capacity)
|
|
if capacity <= 0:
|
|
raise ValueError('capacity must be a positive integer')
|
|
|
|
@property
|
|
def capacity(self):
|
|
return self._capacity
|
|
|
|
@property
|
|
def size(self):
|
|
return len(self._cache)
|
|
|
|
def put(self, key, val):
|
|
if key in self._cache:
|
|
return
|
|
if len(self._cache) >= self.capacity:
|
|
self._cache.popitem(last=False)
|
|
self._cache[key] = val
|
|
|
|
def get(self, key, default=None):
|
|
val = self._cache[key] if key in self._cache else default
|
|
return val
|
|
|
|
|
|
class VideoReader:
|
|
"""Video class with similar usage to a list object.
|
|
|
|
This video warpper class provides convenient apis to access frames.
|
|
There exists an issue of OpenCV's VideoCapture class that jumping to a
|
|
certain frame may be inaccurate. It is fixed in this class by checking
|
|
the position after jumping each time.
|
|
Cache is used when decoding videos. So if the same frame is visited for
|
|
the second time, there is no need to decode again if it is stored in the
|
|
cache.
|
|
|
|
:Example:
|
|
|
|
>>> import annotator.uniformer.mmcv as mmcv
|
|
>>> v = mmcv.VideoReader('sample.mp4')
|
|
>>> len(v) # get the total frame number with `len()`
|
|
120
|
|
>>> for img in v: # v is iterable
|
|
>>> mmcv.imshow(img)
|
|
>>> v[5] # get the 6th frame
|
|
"""
|
|
|
|
def __init__(self, filename, cache_capacity=10):
|
|
|
|
if not filename.startswith(('https://', 'http://')):
|
|
check_file_exist(filename, 'Video file not found: ' + filename)
|
|
self._vcap = cv2.VideoCapture(filename)
|
|
assert cache_capacity > 0
|
|
self._cache = Cache(cache_capacity)
|
|
self._position = 0
|
|
|
|
self._width = int(self._vcap.get(CAP_PROP_FRAME_WIDTH))
|
|
self._height = int(self._vcap.get(CAP_PROP_FRAME_HEIGHT))
|
|
self._fps = self._vcap.get(CAP_PROP_FPS)
|
|
self._frame_cnt = int(self._vcap.get(CAP_PROP_FRAME_COUNT))
|
|
self._fourcc = self._vcap.get(CAP_PROP_FOURCC)
|
|
|
|
@property
|
|
def vcap(self):
|
|
""":obj:`cv2.VideoCapture`: The raw VideoCapture object."""
|
|
return self._vcap
|
|
|
|
@property
|
|
def opened(self):
|
|
"""bool: Indicate whether the video is opened."""
|
|
return self._vcap.isOpened()
|
|
|
|
@property
|
|
def width(self):
|
|
"""int: Width of video frames."""
|
|
return self._width
|
|
|
|
@property
|
|
def height(self):
|
|
"""int: Height of video frames."""
|
|
return self._height
|
|
|
|
@property
|
|
def resolution(self):
|
|
"""tuple: Video resolution (width, height)."""
|
|
return (self._width, self._height)
|
|
|
|
@property
|
|
def fps(self):
|
|
"""float: FPS of the video."""
|
|
return self._fps
|
|
|
|
@property
|
|
def frame_cnt(self):
|
|
"""int: Total frames of the video."""
|
|
return self._frame_cnt
|
|
|
|
@property
|
|
def fourcc(self):
|
|
"""str: "Four character code" of the video."""
|
|
return self._fourcc
|
|
|
|
@property
|
|
def position(self):
|
|
"""int: Current cursor position, indicating frame decoded."""
|
|
return self._position
|
|
|
|
def _get_real_position(self):
|
|
return int(round(self._vcap.get(CAP_PROP_POS_FRAMES)))
|
|
|
|
def _set_real_position(self, frame_id):
|
|
self._vcap.set(CAP_PROP_POS_FRAMES, frame_id)
|
|
pos = self._get_real_position()
|
|
for _ in range(frame_id - pos):
|
|
self._vcap.read()
|
|
self._position = frame_id
|
|
|
|
def read(self):
|
|
"""Read the next frame.
|
|
|
|
If the next frame have been decoded before and in the cache, then
|
|
return it directly, otherwise decode, cache and return it.
|
|
|
|
Returns:
|
|
ndarray or None: Return the frame if successful, otherwise None.
|
|
"""
|
|
|
|
if self._cache:
|
|
img = self._cache.get(self._position)
|
|
if img is not None:
|
|
ret = True
|
|
else:
|
|
if self._position != self._get_real_position():
|
|
self._set_real_position(self._position)
|
|
ret, img = self._vcap.read()
|
|
if ret:
|
|
self._cache.put(self._position, img)
|
|
else:
|
|
ret, img = self._vcap.read()
|
|
if ret:
|
|
self._position += 1
|
|
return img
|
|
|
|
def get_frame(self, frame_id):
|
|
"""Get frame by index.
|
|
|
|
Args:
|
|
frame_id (int): Index of the expected frame, 0-based.
|
|
|
|
Returns:
|
|
ndarray or None: Return the frame if successful, otherwise None.
|
|
"""
|
|
if frame_id < 0 or frame_id >= self._frame_cnt:
|
|
raise IndexError(
|
|
f'"frame_id" must be between 0 and {self._frame_cnt - 1}')
|
|
if frame_id == self._position:
|
|
return self.read()
|
|
if self._cache:
|
|
img = self._cache.get(frame_id)
|
|
if img is not None:
|
|
self._position = frame_id + 1
|
|
return img
|
|
self._set_real_position(frame_id)
|
|
ret, img = self._vcap.read()
|
|
if ret:
|
|
if self._cache:
|
|
self._cache.put(self._position, img)
|
|
self._position += 1
|
|
return img
|
|
|
|
def current_frame(self):
|
|
"""Get the current frame (frame that is just visited).
|
|
|
|
Returns:
|
|
ndarray or None: If the video is fresh, return None, otherwise
|
|
return the frame.
|
|
"""
|
|
if self._position == 0:
|
|
return None
|
|
return self._cache.get(self._position - 1)
|
|
|
|
def cvt2frames(self,
|
|
frame_dir,
|
|
file_start=0,
|
|
filename_tmpl='{:06d}.jpg',
|
|
start=0,
|
|
max_num=0,
|
|
show_progress=True):
|
|
"""Convert a video to frame images.
|
|
|
|
Args:
|
|
frame_dir (str): Output directory to store all the frame images.
|
|
file_start (int): Filenames will start from the specified number.
|
|
filename_tmpl (str): Filename template with the index as the
|
|
placeholder.
|
|
start (int): The starting frame index.
|
|
max_num (int): Maximum number of frames to be written.
|
|
show_progress (bool): Whether to show a progress bar.
|
|
"""
|
|
mkdir_or_exist(frame_dir)
|
|
if max_num == 0:
|
|
task_num = self.frame_cnt - start
|
|
else:
|
|
task_num = min(self.frame_cnt - start, max_num)
|
|
if task_num <= 0:
|
|
raise ValueError('start must be less than total frame number')
|
|
if start > 0:
|
|
self._set_real_position(start)
|
|
|
|
def write_frame(file_idx):
|
|
img = self.read()
|
|
if img is None:
|
|
return
|
|
filename = osp.join(frame_dir, filename_tmpl.format(file_idx))
|
|
cv2.imwrite(filename, img)
|
|
|
|
if show_progress:
|
|
track_progress(write_frame, range(file_start,
|
|
file_start + task_num))
|
|
else:
|
|
for i in range(task_num):
|
|
write_frame(file_start + i)
|
|
|
|
def __len__(self):
|
|
return self.frame_cnt
|
|
|
|
def __getitem__(self, index):
|
|
if isinstance(index, slice):
|
|
return [
|
|
self.get_frame(i)
|
|
for i in range(*index.indices(self.frame_cnt))
|
|
]
|
|
|
|
if index < 0:
|
|
index += self.frame_cnt
|
|
if index < 0:
|
|
raise IndexError('index out of range')
|
|
return self.get_frame(index)
|
|
|
|
def __iter__(self):
|
|
self._set_real_position(0)
|
|
return self
|
|
|
|
def __next__(self):
|
|
img = self.read()
|
|
if img is not None:
|
|
return img
|
|
else:
|
|
raise StopIteration
|
|
|
|
next = __next__
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback):
|
|
self._vcap.release()
|
|
|
|
|
|
def frames2video(frame_dir,
|
|
video_file,
|
|
fps=30,
|
|
fourcc='XVID',
|
|
filename_tmpl='{:06d}.jpg',
|
|
start=0,
|
|
end=0,
|
|
show_progress=True):
|
|
"""Read the frame images from a directory and join them as a video.
|
|
|
|
Args:
|
|
frame_dir (str): The directory containing video frames.
|
|
video_file (str): Output filename.
|
|
fps (float): FPS of the output video.
|
|
fourcc (str): Fourcc of the output video, this should be compatible
|
|
with the output file type.
|
|
filename_tmpl (str): Filename template with the index as the variable.
|
|
start (int): Starting frame index.
|
|
end (int): Ending frame index.
|
|
show_progress (bool): Whether to show a progress bar.
|
|
"""
|
|
if end == 0:
|
|
ext = filename_tmpl.split('.')[-1]
|
|
end = len([name for name in scandir(frame_dir, ext)])
|
|
first_file = osp.join(frame_dir, filename_tmpl.format(start))
|
|
check_file_exist(first_file, 'The start frame not found: ' + first_file)
|
|
img = cv2.imread(first_file)
|
|
height, width = img.shape[:2]
|
|
resolution = (width, height)
|
|
vwriter = cv2.VideoWriter(video_file, VideoWriter_fourcc(*fourcc), fps,
|
|
resolution)
|
|
|
|
def write_frame(file_idx):
|
|
filename = osp.join(frame_dir, filename_tmpl.format(file_idx))
|
|
img = cv2.imread(filename)
|
|
vwriter.write(img)
|
|
|
|
if show_progress:
|
|
track_progress(write_frame, range(start, end))
|
|
else:
|
|
for i in range(start, end):
|
|
write_frame(i)
|
|
vwriter.release()
|
|
|
