|
|
|
|
|
import ctypes
|
|
|
import inspect
|
|
|
import os
|
|
|
import os.path as osp
|
|
|
import random
|
|
|
import string
|
|
|
from types import FunctionType, ModuleType
|
|
|
from typing import Optional, Union
|
|
|
|
|
|
import cv2
|
|
|
import mmcv
|
|
|
import numpy as np
|
|
|
|
|
|
|
|
|
def get_random_string(length: int = 15) -> str:
|
|
|
"""Get random string with letters and digits.
|
|
|
|
|
|
Args:
|
|
|
length (int): Length of random string. Defaults to 15.
|
|
|
"""
|
|
|
return ''.join(
|
|
|
random.choice(string.ascii_letters + string.digits)
|
|
|
for _ in range(length))
|
|
|
|
|
|
|
|
|
def get_thread_id() -> int:
|
|
|
"""Get current thread id."""
|
|
|
|
|
|
thread_id = ctypes.CDLL('libc.so.6').syscall(186)
|
|
|
return thread_id
|
|
|
|
|
|
|
|
|
def get_shm_dir() -> str:
|
|
|
"""Get shm dir for temporary usage."""
|
|
|
return '/dev/shm'
|
|
|
|
|
|
|
|
|
def frame_extract(video_path: str,
|
|
|
short_side: Optional[int] = None,
|
|
|
out_dir: str = './tmp'):
|
|
|
"""Extract frames given video_path.
|
|
|
|
|
|
Args:
|
|
|
video_path (str): The video path.
|
|
|
short_side (int): Target short-side of the output image.
|
|
|
Defaults to None, means keeping original shape.
|
|
|
out_dir (str): The output directory. Defaults to ``'./tmp'``.
|
|
|
"""
|
|
|
|
|
|
target_dir = osp.join(out_dir, osp.basename(osp.splitext(video_path)[0]))
|
|
|
os.makedirs(target_dir, exist_ok=True)
|
|
|
|
|
|
frame_tmpl = osp.join(target_dir, 'img_{:06d}.jpg')
|
|
|
assert osp.exists(video_path), f'file not exit {video_path}'
|
|
|
vid = cv2.VideoCapture(video_path)
|
|
|
frames = []
|
|
|
frame_paths = []
|
|
|
flag, frame = vid.read()
|
|
|
cnt = 0
|
|
|
new_h, new_w = None, None
|
|
|
while flag:
|
|
|
if short_side is not None:
|
|
|
if new_h is None:
|
|
|
h, w, _ = frame.shape
|
|
|
new_w, new_h = mmcv.rescale_size((w, h), (short_side, np.Inf))
|
|
|
frame = mmcv.imresize(frame, (new_w, new_h))
|
|
|
|
|
|
frames.append(frame)
|
|
|
frame_path = frame_tmpl.format(cnt + 1)
|
|
|
frame_paths.append(frame_path)
|
|
|
|
|
|
cv2.imwrite(frame_path, frame)
|
|
|
cnt += 1
|
|
|
flag, frame = vid.read()
|
|
|
|
|
|
return frame_paths, frames
|
|
|
|
|
|
|
|
|
class VideoWriter():
|
|
|
|
|
|
def __init__(self, video_file, fps):
|
|
|
self.video_file = video_file
|
|
|
self.fps = fps
|
|
|
if video_file.endswith('.mp4'):
|
|
|
self.fourcc = 'mp4v'
|
|
|
elif video_file.endswith('.avi'):
|
|
|
self.fourcc = 'XVID'
|
|
|
|
|
|
out_dir = osp.dirname(osp.abspath(self.video_file))
|
|
|
if not osp.exists(out_dir):
|
|
|
os.makedirs(out_dir, exist_ok=True)
|
|
|
|
|
|
def _init_cv2_writer(self, frame):
|
|
|
from cv2 import VideoWriter, VideoWriter_fourcc
|
|
|
height, width = frame.shape[:2]
|
|
|
resolution = (width, height)
|
|
|
self.writer = VideoWriter(self.video_file,
|
|
|
VideoWriter_fourcc(*self.fourcc), self.fps,
|
|
|
resolution)
|
|
|
|
|
|
def write_frame(self, frame):
|
|
|
if not getattr(self, 'writer', None):
|
|
|
self._init_cv2_writer(frame)
|
|
|
self.writer.write(frame)
|
|
|
|
|
|
def release(self):
|
|
|
self.writer.release()
|
|
|
|
|
|
def __enter__(self):
|
|
|
return self
|
|
|
|
|
|
def __exit__(self, type, value, trace):
|
|
|
self.release()
|
|
|
|
|
|
|
|
|
def get_str_type(module: Union[str, ModuleType, FunctionType]) -> str:
|
|
|
"""Return the string type name of module.
|
|
|
|
|
|
Args:
|
|
|
module (str | ModuleType | FunctionType):
|
|
|
The target module class
|
|
|
|
|
|
Returns:
|
|
|
Class name of the module
|
|
|
"""
|
|
|
if isinstance(module, str):
|
|
|
str_type = module
|
|
|
elif inspect.isclass(module) or inspect.isfunction(module):
|
|
|
str_type = module.__name__
|
|
|
else:
|
|
|
return None
|
|
|
|
|
|
return str_type
|
|
|
|
