1#!/usr/bin/env python3
2#
3#   Copyright 2019 - The Android Open Source Project
4#
5#   Licensed under the Apache License, Version 2.0 (the "License");
6#   you may not use this file except in compliance with the License.
7#   You may obtain a copy of the License at
8#
9#       http://www.apache.org/licenses/LICENSE-2.0
10#
11#   Unless required by applicable law or agreed to in writing, software
12#   distributed under the License is distributed on an "AS IS" BASIS,
13#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14#   See the License for the specific language governing permissions and
15#   limitations under the License.
16
17import logging
18import os
19import pyaudio
20import wave
21
22from acts import context
23
24
25WAVE_FILE_TEMPLATE = 'recorded_audio_%s.wav'
26ADB_PATH = 'sdcard/Music/'
27ADB_FILE = 'rec.pcm'
28
29
30class AudioCaptureBase(object):
31    """Base class for Audio capture."""
32
33    def __init__(self):
34
35        self.wave_file = os.path.join(self.log_path, WAVE_FILE_TEMPLATE)
36        self.file_dir = self.log_path
37
38    @property
39    def log_path(self):
40        """Returns current log path."""
41        current_context = context.get_current_context()
42        full_out_dir = os.path.join(current_context.get_full_output_path(),
43                                    'AudioCapture')
44
45        os.makedirs(full_out_dir, exist_ok=True)
46        return full_out_dir
47
48    @property
49    def next_fileno(self):
50        counter = 0
51        while os.path.exists(self.wave_file % counter):
52            counter += 1
53        return counter
54
55    @property
56    def last_fileno(self):
57        return self.next_fileno - 1
58
59    @property
60    def get_last_record_duration_millis(self):
61        """Get duration of most recently recorded file.
62
63        Returns:
64            duration (float): duration of recorded file in milliseconds.
65        """
66        latest_file_path = self.wave_file % self.last_fileno
67        print (latest_file_path)
68        with wave.open(latest_file_path, 'r') as f:
69            frames = f.getnframes()
70            rate = f.getframerate()
71            duration = (frames / float(rate)) * 1000
72        return duration
73
74    def write_record_file(self, audio_params, frames):
75        """Writes the recorded audio into the file.
76
77        Args:
78            audio_params: A dict with audio configuration.
79            frames: Recorded audio frames.
80
81        Returns:
82            file_name: wave file name.
83        """
84        file_name = self.wave_file % self.next_fileno
85        logging.debug('writing to %s' % file_name)
86        wf = wave.open(file_name, 'wb')
87        wf.setnchannels(audio_params['channel'])
88        wf.setsampwidth(audio_params['sample_width'])
89        wf.setframerate(audio_params['sample_rate'])
90        wf.writeframes(frames)
91        wf.close()
92        return file_name
93
94
95class CaptureAudioOverAdb(AudioCaptureBase):
96    """Class to capture audio over android device which acts as the
97    a2dp sink or hfp client. This captures the digital audio and converts
98    to analog audio for post processing.
99    """
100
101    def __init__(self, ad, audio_params):
102        """Initializes CaptureAudioOverAdb.
103
104        Args:
105            ad: An android device object.
106            audio_params: Dict containing audio record settings.
107        """
108        super().__init__()
109        self._ad = ad
110        self.audio_params = audio_params
111        self.adb_path = None
112
113    def start(self):
114        """Start the audio capture over adb."""
115        self.adb_path = os.path.join(ADB_PATH, ADB_FILE)
116        cmd = 'ap2f --usage 1 --start --duration {} --target {}'.format(
117            self.audio_params['duration'], self.adb_path,
118        )
119        self._ad.adb.shell_nb(cmd)
120
121    def stop(self):
122        """Stops the audio capture and stores it in wave file.
123
124        Returns:
125            File name of the recorded file.
126        """
127        cmd = '{} {}'.format(self.adb_path, self.file_dir)
128        self._ad.adb.pull(cmd)
129        self._ad.adb.shell('rm {}'.format(self.adb_path))
130        return self._convert_pcm_to_wav()
131
132    def _convert_pcm_to_wav(self):
133        """Converts raw pcm data into wave file.
134
135        Returns:
136            file_path: Returns the file path of the converted file
137            (digital to analog).
138        """
139        file_to_read = os.path.join(self.file_dir, ADB_FILE)
140        with open(file_to_read, 'rb') as pcm_file:
141            frames = pcm_file.read()
142        file_path = self.write_record_file(self.audio_params, frames)
143        return file_path
144
145
146class CaptureAudioOverLocal(AudioCaptureBase):
147    """Class to capture audio on local server using the audio input devices
148    such as iMic/AudioBox. This class mandates input deivce to be connected to
149    the machine.
150    """
151    def __init__(self, audio_params):
152        """Initializes CaptureAudioOverLocal.
153
154        Args:
155            audio_params: Dict containing audio record settings.
156        """
157        super().__init__()
158        self.audio_params = audio_params
159        self.channels = self.audio_params['channel']
160        self.chunk = self.audio_params['chunk']
161        self.sample_rate = self.audio_params['sample_rate']
162        self.__input_device = None
163        self.audio = None
164        self.frames = []
165
166    @property
167    def name(self):
168        return self.__input_device["name"]
169
170    def __get_input_device(self):
171        """Checks for the audio capture device."""
172        if self.__input_device is None:
173            for i in range(self.audio.get_device_count()):
174                device_info = self.audio.get_device_info_by_index(i)
175                logging.debug('Device Information: {}'.format(device_info))
176                if self.audio_params['input_device'] in device_info['name']:
177                    self.__input_device = device_info
178                    break
179            else:
180                raise DeviceNotFound(
181                    'Audio Capture device {} not found.'.format(
182                        self.audio_params['input_device']))
183        return self.__input_device
184
185    def start(self, trim_beginning=0, trim_end=0):
186        """Starts audio recording on host machine.
187
188        Args:
189            trim_beginning: how many seconds to trim from the beginning
190            trim_end: how many seconds to trim from the end
191        """
192        self.audio = pyaudio.PyAudio()
193        self.__input_device = self.__get_input_device()
194        stream = self.audio.open(
195            format=pyaudio.paInt16,
196            channels=self.channels,
197            rate=self.sample_rate,
198            input=True,
199            frames_per_buffer=self.chunk,
200            input_device_index=self.__input_device['index'])
201        b_chunks = trim_beginning * (self.sample_rate // self.chunk)
202        e_chunks = trim_end * (self.sample_rate // self.chunk)
203        total_chunks = self.sample_rate // self.chunk * self.audio_params[
204            'duration']
205        for i in range(total_chunks):
206            try:
207                data = stream.read(self.chunk, exception_on_overflow=False)
208            except IOError as ex:
209                logging.error('Cannot record audio: {}'.format(ex))
210                return False
211            if b_chunks <= i < total_chunks - e_chunks:
212                self.frames.append(data)
213
214        stream.stop_stream()
215        stream.close()
216
217    def stop(self):
218        """Terminates the pulse audio instance.
219
220        Returns:
221            File name of the recorded audio file.
222        """
223        self.audio.terminate()
224        frames = b''.join(self.frames)
225        return self.write_record_file(self.audio_params, frames)
226
227
228class DeviceNotFound(Exception):
229    """Raises exception if audio capture device is not found."""
230