Coverage for ibllib/io/ffmpeg.py: 97%

32 statements  

« prev     ^ index     » next       coverage.py v7.5.4, created at 2024-07-08 17:16 +0100

1from pathlib import Path 

2import subprocess 

3import logging 

4 

5from ibllib.io.video import get_video_meta 

6 

7_logger = logging.getLogger(__name__) 

8 

9 

10def compress(file_in, file_out, command, remove_original=True): 

11 """ 

12 runs a command of the form 'ffmpeg -i {file_in} -c:a flac -nostats {file_out}' 

13 using a supbprocess 

14 audio compression for ephys: 

15 `"ffmpeg -i {file_in} -c:a flac -nostats {file_out}"` 

16 video compression for ephys: 

17 `"('ffmpeg -i {file_in} -codec:v libx264 -preset slow -crf 17 ' 

18 '-nostats -loglevel 0 -codec:a copy {file_out}')"` 

19 

20 :param file_in: full file path of input 

21 :param file_out: full file path of output 

22 :param command: string ready to be formatted with `file_in` and `file_out` 

23 :param remove_original: if true, remove uncompressed file 

24 return: file_out 

25 """ 

26 

27 file_in = Path(file_in) 1ahcdebfg

28 file_out = Path(file_out) 1ahcdebfg

29 # if the output file already exists, overwrite it 

30 if file_out.exists(): 1ahcdebfg

31 file_out.unlink() 

32 command2run = command.format(file_in=str(file_in), file_out=str(file_out)) 1ahcdebfg

33 process = subprocess.Popen(command2run, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) 1ahcdebfg

34 info, error = process.communicate() 1ahcdebfg

35 assert process.returncode == 0, f'compression failed for {file_in}: {error}' 1ahcdebfg

36 # if the command was successful delete the original file 

37 if remove_original: 1ahcdebfg

38 file_in.unlink() 1hb

39 return process.returncode, file_out 1ahcdebfg

40 

41 

42def iblrig_video_compression(session_path, command, verify_output=True): 

43 """ 

44 Compresses avi files from the raw_video_data folder into mp4. Once the compression completes 

45 the avi file is removed. If mp4 already exist, the funciton returns their file path 

46 :param session_path: the session directory path 

47 :param command: string ready to be formatted with `file_in` and `file_out` 

48 :param verify_output: if true assert that output video is readable 

49 for ephys: 

50 >>> command = ('ffmpeg -i {file_in} -y -codec:v libx264 -preset slow -crf 17 ' 

51 >>> '-nostats -loglevel 0 -codec:a copy {file_out}') 

52 for training: 

53 >>> command = ('ffmpeg -i {file_in} -y -codec:v libx264 -preset slow -crf 29 ' 

54 >>> '-nostats -loglevel 0 -codec:a copy {file_out}') 

55 :return: list of compressed files 

56 """ 

57 output_files = list(session_path.joinpath("raw_video_data").rglob('_iblrig_*.mp4')) 1acdeijbfg

58 rig_avi_files = list(session_path.joinpath("raw_video_data").rglob('_iblrig_*.avi')) 1acdeijbfg

59 # first compress everything (the rationale is not to delete anything if there is a crash) 

60 for file_in in rig_avi_files: 1acdeijbfg

61 _logger.info(f"compressing {file_in}") 1acdebfg

62 file_out = file_in.with_suffix('.mp4') 1acdebfg

63 status, fout = compress(file_in=file_in, file_out=file_out, 1acdebfg

64 command=command, remove_original=False) 

65 output_files.append(fout) 1acdebfg

66 

67 if verify_output: 1acdeijbfg

68 for file_out in output_files: 1acdeijbfg

69 meta = get_video_meta(file_out) 1acdeijbfg

70 assert meta.length > 0 and meta.size > 0, f'Video file empty: {file_out}' 1acdeijbfg

71 

72 # then remove everything 

73 for file_in in rig_avi_files: 1acdeijbfg

74 file_in.unlink() 1acdebfg

75 return output_files 1acdeijbfg