Shortcuts

Source code for mmaction.utils.misc

# Copyright (c) OpenMMLab. All rights reserved.
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


[docs]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))
[docs]def get_thread_id() -> int: """Get current thread id.""" # use ctype to find thread id thread_id = ctypes.CDLL('libc.so.6').syscall(186) return thread_id
[docs]def get_shm_dir() -> str: """Get shm dir for temporary usage.""" return '/dev/shm'
[docs]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'``. """ # Load the video, extract frames into OUT_DIR/video_name target_dir = osp.join(out_dir, osp.basename(osp.splitext(video_path)[0])) os.makedirs(target_dir, exist_ok=True) # Should be able to handle videos up to several hours 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()
[docs]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