Coverage for ibllib/pipes/audio_tasks.py: 97%

37 statements  

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

1import logging 

2 

3from ibllib.pipes import base_tasks 

4from ibllib.io import ffmpeg 

5from ibllib.io.extractors import training_audio 

6 

7_logger = logging.getLogger('ibllib') 

8 

9 

10class AudioCompress(base_tasks.AudioTask): 

11 

12 cpu = 2 

13 priority = 10 

14 job_size = 'small' 

15 

16 @property 

17 def signature(self): 

18 signature = { 1b

19 'input_files': [('_iblrig_micData.raw.wav', self.device_collection, True)], 

20 'output_files': [('_iblrig_micData.raw.flac', self.device_collection, True)] 

21 } 

22 return signature 1b

23 

24 def _run(self, overwrite=False): 

25 

26 command = "ffmpeg -i {file_in} -y -nostdin -c:a flac -nostats {file_out}" 1b

27 file_in = next(self.session_path.joinpath(self.device_collection).rglob("_iblrig_micData.raw.wav"), None) 1b

28 if file_in is None: 1b

29 return 

30 file_out = file_in.with_suffix(".flac") 1b

31 status, output_file = ffmpeg.compress(file_in=file_in, file_out=file_out, command=command) 1b

32 return [output_file] 1b

33 

34 

35class AudioSync(base_tasks.AudioTask): 

36 """ 

37 Extracts audio events and sync. N.B currently only supports bpod with xonar sound system 

38 """ 

39 

40 cpu = 2 

41 priority = 10 

42 job_size = 'small' 

43 

44 def __init__(self, session_path, **kwargs): 

45 super().__init__(session_path, **kwargs) 1dceghijf

46 # Task collection (this needs to be specified in the task kwargs) 

47 self.collection = self.get_task_collection(kwargs.get('collection', None)) 1dceghijf

48 

49 @property 

50 def signature(self): 

51 signature = { 1dcef

52 'input_files': [('_iblrig_micData.raw.wav', self.device_collection, True)], 

53 'output_files': [('_iblmic_audioOnsetGoCue.times_mic.npy', self.device_collection, True), 

54 ('_iblmic_audioSpectrogram.frequencies.npy', self.device_collection, True), 

55 ('_iblmic_audioSpectrogram.power.npy', self.device_collection, True), 

56 ('_iblmic_audioSpectrogram.times_mic.npy', self.device_collection, True)] 

57 } 

58 return signature 1dcef

59 

60 def _run(self): 

61 if self.sync == 'bpod': 1dcef

62 return training_audio.extract_sound(self.session_path, task_collection=self.collection, 1def

63 device_collection=self.device_collection, save=True, delete=True) 

64 else: 

65 _logger.warning('Audio Syncing not yet implemented for FPGA') 1c

66 return 1c