Coverage for /home/ubuntu/cv-depot/python/cv_depot/core/video.py: 100%

1from typing import Any, Dict, Union # noqa F401 

2from cv_depot.core.types import Filepath # noqa F401 

3 

4import os 

5from pathlib import Path 

6import re 

7 

8import ffmpeg 

9from lunchbox.enforce import Enforce 

10 

11from cv_depot.core.enum import VideoCodec, VideoFormat 

12# ------------------------------------------------------------------------------ 

13 

14 

15def get_video_metadata(filepath): 

16 # type: (Filepath) -> Dict[str, Any] 

17 ''' 

18 Retrieve video file metadata with ffmpeg. 

19 

20 Args: 

21 filepath (str or Path): Path to video file. 

22 

23 Raises: 

24 EnforceError: If filepath is not a file or does not exist. 

25 EnforceError: If filepath format is unsupported. 

26 

27 Returns: 

28 dict: Metadata. 

29 ''' 

30 filepath = Path(filepath) 

31 

32 msg = f'{filepath} is not a file or does not exist.' 

33 Enforce(filepath.is_file(), '==', True, message=msg) 

34 

35 ext = filepath.suffix[1:] 

36 formats = ['m4v', 'mov', 'mp4', 'mpg', 'mpeg'] 

37 msg = 'Illegal format: {a}. Legal formats: {b}.' 

38 Enforce(ext.lower(), 'in', formats, message=msg) 

39 # -------------------------------------------------------------------------- 

40 

41 meta = ffmpeg.probe(filepath.as_posix())['streams'][0] 

42 lut = {'16': 'FLOAT16', '32': 'FLOAT32', '8': 'UINT8'} 

43 output = dict( 

44 width=meta['width'], 

45 height=meta['height'], 

46 num_frames=int(meta['nb_frames']), 

47 frame_rate=int(meta['r_frame_rate'].split('/')[0]), 

48 bit_depth=lut[meta['bits_per_raw_sample']], 

49 codec=meta['codec_name'], 

50 ) 

51 return output 

52 

53 

54def write_video( 

55 source, # type: Filepath 

56 target, # type: Filepath 

57 framerate=24, # type: int 

58 codec=VideoCodec.H264, # type: VideoCodec 

59 format_=VideoFormat.MP4 # type: VideoFormat 

60): 

61 # type: (...) -> None 

62 ''' 

63 Writes given input file or file pattern to given target file. 

64 

65 Args: 

66 source (str or Path): Source image filepath or file pattern. 

67 target (str or Path): Target file. 

68 framerate (int, optional): Video framerate. Default: 24 fps. 

69 codec (VideoCodec, optional): Video codec. Default: VideoCodec.H264. 

70 format_ (VideoFormat, optional): Video container format. 

71 Default: VideoFormat.MP4 

72 

73 Raises: 

74 EnforceError: If source is not a filepath or file pattern. 

75 EnforceError: framerate is not an integer greater than 0. 

76 EnforceError: If codec is illegal. 

77 EnforceError: If format is illegal. 

78 ''' 

79 source = Path(source).as_posix() 

80 target = Path(target).as_posix() 

81 if not Path(source).is_file(): 

82 msg = 'Source must be a video file or file pattern with {a} in it. ' 

83 msg += 'Given value: {b}.' 

84 Enforce('{frame}', 'in', source, message=msg) 

85 source = re.sub(r'\{frame\}', '%04d', source) 

86 

87 msg = 'Framerate must be an integer greater than 0. Given value: {a}.' 

88 Enforce(framerate, 'instance of', int, message=msg) 

89 Enforce(framerate, '>', 0, message=msg) 

90 

91 Enforce(codec, 'instance of', VideoCodec) 

92 Enforce(format_, 'instance of', VideoFormat) 

93 # -------------------------------------------------------------------------- 

94 

95 os.makedirs(Path(target).parent, exist_ok=True) 

96 ffmpeg \ 

97 .input(source, framerate=framerate) \ 

98 .output(target, vcodec=codec.ffmpeg_code, format=format_.extension) \ 

99 .run(overwrite_output=True)