1#!/usr/bin/env python3

2# This file is part of Checkbox.

3#

4# Copyright 2023 Canonical Ltd.

5# Written by:

6# Hanhsuan Lee <hanhsuan.lee@canonical.com>

7#

8# Checkbox is free software: you can redistribute it and/or modify

9# it under the terms of the GNU General Public License version 3,

10# as published by the Free Software Foundation.

11#

12# Checkbox is distributed in the hope that it will be useful,

13# but WITHOUT ANY WARRANTY; without even the implied warranty of

14# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the

15# GNU General Public License for more details.

16#

17# You should have received a copy of the GNU General Public License

18# along with Checkbox. If not, see <http://www.gnu.org/licenses/>.

19

20 from __future__ import division, print_function

21 import argparse

22 import collections

23 import json

24 import logging

25 import math

26 import re

27 import subprocess

28 import sys

29 import time

30

31 try:

32 import gi

33

34 gi.require_version("GLib", "2.0")

35 gi.require_version("Gst", "1.0")

36 from gi.repository import GObject

37 from gi.repository import Gst

38 from gi.repository import GLib

39

40 Gst.init(None) # This has to be done very early so it can find elements

41 except ImportError:

42 print(

43 "Can't import module: %s. it may not be available for this"

44 "version of Python, which is: " % sys.exc_info()[1],

45 file=sys.stderr,

46 )

47 print((sys.version), file=sys.stderr)

48 sys.exit(127)

49

50

51# Frequency bands for FFT

52 BINS = 256

53# How often to take a sample and do FFT on it.

54 FFT_INTERVAL = 100000000 # In nanoseconds, so this is every 1/10th second

55# Sampling frequency. The effective maximum frequency we can analyze is

56# half of this (see Nyquist's theorem)

57 SAMPLING_FREQUENCY = 44100

58# The default test frequency is in the middle of the band that contains 5000Hz

59# This frequency was determined experimentally to be high enough but more

60# reliable than others we tried.

61 DEFAULT_TEST_FREQUENCY = 5035

62# only sample a signal when peak level is in this range (in dB attenuation,

63# 0 means no attenuation (and horrible clipping).

64 REC_LEVEL_RANGE = (-2.0, -12.0)

65# For our test signal to be considered present, it has to be this much higher

66# than the base level (minimum magnitude). This is in dB.

67 MAGNITUDE_THRESHOLD = 2.5

68# Volume for the sample tone (in %)

69 PLAY_VOLUME = 70

70

71

72 class PIDController(object):

73 """A Proportional-Integrative-Derivative controller (PID) controls a

74 process's output to try to maintain a desired output value (known as

75 'setpoint', by continually adjusting the process's input.

76

77 It does so by calculating the "error" (difference between output and

78 setpoint) and attempting to minimize it manipulating the input.

79

80 The desired change to the input is calculated based on error and three

81 constants (kp, ki and kd). These values can be interpreted in terms of

82 time: P depends on the present error, I on the accumulation of past errors,

83 and D is a prediction of future errors, based on current rate of change.

84 The weighted sum of these three actions is used to adjust the process via a

85 control element.

86

87 In practice, kp, ki and kd are process-dependent and usually have to

88 be tweaked by hand, but once reasonable constants are arrived at, they

89 can apply to a particular process without further modification.

90

91 """

92

93 def __init__(self, kp, ki, kd, setpoint=0):

94 """Creates a PID controller with given constants and setpoint.

95

96 Arguments:

97 kp, ki, kd: PID constants, see class description.

98 setpoint: desired output value; calls to input_change with

99 a process output reading will return a desired change

100 to the input to attempt matching output to this value.

101 """

102 self.setpoint = setpoint

103 self.kp = kp

104 self.ki = ki

105 self.kd = kd

106 self._integral = 0

107 self._previous_error = 0

108 self._change_limit = 0

109

110 def input_change(self, process_feedback, dt):

111 """Calculates desired input value change.

112

113 Based on process feedback and time interval (dt).

114 """

115 error = self.setpoint - process_feedback

116 self._integral = self._integral + (error * dt)

117 derivative = (error - self._previous_error) / dt

118 self._previous_error = error

119 input_change = (

120 (self.kp * error)

121 + (self.ki * self._integral)

122 + (self.kd * derivative)

123 )

124 if self._change_limit and abs(input_change) > abs(self._change_limit):

125 sign = input_change / abs(input_change)

126 input_change = sign * self._change_limit

127 return input_change

128

129 def set_change_limit(self, limit):

130 """Ensures that input value changes are lower than limit.

131

132 Setting limit of zero disables this.

133 """

134 self._change_limit = limit

135

136

137 class VolumeController(object):

138 pw_types = {"input": "SOURCE", "output": "SINK"}

139

140 def __init__(self, type, logger=None):

141 """Initializes the volume controller.

142

143 Arguments:

144 type: either input or output

145

146 """

147 self.type = type

148 self._volume = None

149 self.identifier = None

150 self.logger = logger

151

152 def set_volume(self, volume):

153 if not 0 <= volume <= 100:

154 return False

155 command = [

156 "wpctl",

157 "set-volume",

158 "@DEFAULT_AUDIO_%s@" % (self.pw_types[self.type]),

159 str(int(volume)) + "%",

160 ]

161 self._wpctl_output(command)

162 self._volume = volume

163 return True

164

165 def get_volume(self):

166 return self._volume

167

168 def mute(self, mute):

169 mute = str(int(mute))

170 command = [

171 "wpctl",

172 "set-mute",

173 "@DEFAULT_AUDIO_%s@" % (self.pw_types[self.type]),

174 mute,

175 ]

176 self._wpctl_output(command)

177 return True

178

179 def _wpctl_output(self, command):

180 # This method mainly calls wpctl (hence the name). Since wpctl may

181 # return a failure if the audio layer is not yet initialized, we will

182 # try running a few times in case of failure. All our invocations of

183 # wpctl should be "idempotent" so repeating them should not have

184 # any bad effects.

185 for attempt in range(0, 3):

186 try:

187 return subprocess.check_output(

188 command, universal_newlines=True

189 )

190 except subprocess.CalledProcessError:

191 time.sleep(5)

192 except FileNotFoundError:

193 break

194 self.logger.error("Fail to execute: {}".format(command))

195 sys.exit(1)

196

197

198 class FileDumper(object):

199 def write_to_file(self, filename, data):

200 try:

201 with open(filename, "w") as f:

202 for i in data:

203 print(i, file=f)

204 return_value = True

205 except (TypeError, IOError) as e:

206 logging.error(repr(e))

207 return_value = False

208 return return_value

209

210

211 class SpectrumAnalyzer(object):

212 def __init__(self, points, sampling_frequency=44100, wanted_samples=50):

213 self.spectrum = [0] * points

214 self.number_of_samples = 0

215 self.wanted_samples = wanted_samples

216 self.sampling_frequency = sampling_frequency

217 # Frequencies should contain *real* frequency which is half of

218 # the sampling frequency

219 self.frequencies = [

220 ((sampling_frequency / 2.0) / points) * i for i in range(points)

221 ]

222

223 def _average(self):

224 return sum(self.spectrum) / len(self.spectrum)

225

226 def sample(self, sample):

227 if len(sample) != len(self.spectrum):

228 return

229 self.spectrum = [

230 ((old * self.number_of_samples) + new)

231 / (self.number_of_samples + 1)

232 for old, new in zip(self.spectrum, sample)

233 ]

234 self.number_of_samples += 1

235

236 def frequencies_with_peak_magnitude(self, threshold=1.0):

237 # First establish the base level

238 per_magnitude_bins = collections.defaultdict(int)

239 for magnitude in self.spectrum:

240 per_magnitude_bins[magnitude] += 1

241 base_level = max(

242 per_magnitude_bins, key=lambda x: per_magnitude_bins[x]

243 )

244 # Now return all values that are higher (more positive)

245 # than base_level + threshold

246 peaks = []

247 for i in range(1, len(self.spectrum) - 1):

248 first_index = i - 1

249 last_index = i + 1

250 if ( 250 ↛ 255 line 250 didn't jump to line 255 because the condition on line 250 was never true

251 self.spectrum[first_index] < self.spectrum[i]

252 and self.spectrum[last_index] < self.spectrum[i]

253 and self.spectrum[i] > base_level + threshold

254 ):

255 peaks.append(i)

256

257 return peaks

258

259 def frequency_band_for(self, frequency):

260 """Convenience function to tell me which band

261 a frequency is contained in

262 """

263 # Note that actual frequencies are half of what the sampling

264 # frequency would tell us. If SF is 44100 then maximum actual

265 # frequency is 22050, and if I have 10 frequency bins each will

266 # contain only 2205 Hz, not 4410 Hz.

267 max_frequency = self.sampling_frequency / 2

268 if frequency > max_frequency or frequency < 0:

269 return None

270 band = float(frequency) / (max_frequency / len(self.spectrum))

271 return int(math.ceil(band)) - 1

272

273 def frequencies_for_band(self, band):

274 """Convenience function to tell me the delimiting frequencies

275 for a band

276 """

277 if band >= len(self.spectrum) or band < 0:

278 return None

279 lower = self.frequencies[band]

280 upper = lower + ((self.sampling_frequency / 2.0) / len(self.spectrum))

281 return (lower, upper)

282

283 def sampling_complete(self):

284 return self.number_of_samples >= self.wanted_samples

285

286

287 class GStreamerMessageHandler(object):

288 def __init__(

289 self,

290 rec_level_range,

291 logger,

292 volumecontroller,

293 pidcontroller,

294 spectrum_analyzer,

295 ):

296 """Initializes the message handler. It knows how to handle

297 spectrum and level gstreamer messages.

298

299 Arguments:

300 rec_level_range: tuple with acceptable recording level

301 ranges

302 logger: logging object with debug, info, error methods.

303 volumecontroller: an instance of VolumeController to use

304 to adjust RECORDING level

305 pidcontroller: a PID controller instance which helps control

306 volume

307 spectrum_analyzer: instance of SpectrumAnalyzer to collect

308 data from spectrum messages

309

310 """

311 self.current_level = sys.maxsize

312 self.logger = logger

313 self.pid_controller = pidcontroller

314 self.rec_level_range = rec_level_range

315 self.spectrum_analyzer = spectrum_analyzer

316 self.volume_controller = volumecontroller

317

318 def set_quit_method(self, method):

319 """Method that will be called when sampling is complete."""

320 self._quit_method = method

321

322 def bus_message_handler(self, bus, message):

323 if message.type == Gst.MessageType.ELEMENT: 323 ↛ exit line 323 didn't return from function 'bus_message_handler' because the condition on line 323 was always true

324 message_name = message.get_structure().get_name()

325 if message_name == "spectrum":

326 # TODO: Due to an upstream bug, a structure's get_value method

327 # doesn't work if the value in question is an array (as is the

328 # case with the magnitudes).

329 # https://bugzilla.gnome.org/show_bug.cgi?id=693168

330 # We have to resort to parsing the string representation of the

331 # structure. It's an ugly hack but it works.

332 # Ideally we'd be able to say this to get fft_magnitudes:

333 # message.get_structure.get_value('magnitude').

334 # If an upstream fix ever makes it into gstreamer,

335 # remember to remove this hack and the parse_spectrum method

336 struct_string = message.get_structure().to_string()

337 structure = parse_spectrum_message_structure(struct_string)

338 fft_magnitudes = structure["magnitude"]

339 self.spectrum_method(self.spectrum_analyzer, fft_magnitudes)

340

341 if message_name == "level":

342 # peak_value is our process feedback

343 # It's returned as an array, so I need the first (and only)

344 # element

345 peak_value = message.get_structure().get_value("peak")[0]

346 self.level_method(

347 peak_value, self.pid_controller, self.volume_controller

348 )

349

350 # Adjust recording level

351 def level_method(self, level, pid_controller, volume_controller):

352 # If volume controller doesn't return a valid volume,

353 # we can't control it :(

354 current_volume = volume_controller.get_volume()

355 if current_volume is None:

356 self.logger.error(

357 "Unable to control recording volume. "

358 "Test results may be wrong"

359 )

360 return

361 self.current_level = level

362 change = pid_controller.input_change(level, 0.10)

363 if self.logger: 363 ↛ 373 line 363 didn't jump to line 373 because the condition on line 363 was always true

364 self.logger.debug(

365 "Peak level: %(peak_level).2f, "

366 "volume: %(volume)d%%, Volume change: %(change)f%%"

367 % {

368 "peak_level": level,

369 "change": change,

370 "volume": current_volume,

371 }

372 )

373 volume_controller.set_volume(current_volume + change)

374

375 # Only sample if level is within the threshold

376 def spectrum_method(self, analyzer, spectrum):

377 if ( 377 ↛ 385 line 377 didn't jump to line 385 because the condition on line 377 was always true

378 self.rec_level_range[1] <= self.current_level

379 or self.current_level <= self.rec_level_range[0]

380 ):

381 self.logger.debug(

382 "Sampling, recorded %d samples" % analyzer.number_of_samples

383 )

384 analyzer.sample(spectrum)

385 if analyzer.sampling_complete() and self._quit_method: 385 ↛ 386 line 385 didn't jump to line 386 because the condition on line 385 was never true

386 self.logger.info("Sampling complete, ending process")

387 self._quit_method()

388

389

390 class GstAudioObject(object):

391 def __init__(self):

392 self.class_name = self.__class__.__name__

393

394 def _set_state(self, state, description):

395 self.pipeline.set_state(state)

396 message = "%s: %s" % (self.class_name, description)

397 if self.logger: 397 ↛ exit line 397 didn't return from function '_set_state' because the condition on line 397 was always true

398 self.logger.info(message)

399

400 def start(self):

401 self._set_state(Gst.State.PLAYING, "Starting")

402

403 def stop(self):

404 self._set_state(Gst.State.NULL, "Stopping")

405

406

407 class Player(GstAudioObject):

408 def __init__(self, frequency=DEFAULT_TEST_FREQUENCY, logger=None):

409 super(Player, self).__init__()

410 self.pipeline_description = (

411 "audiotestsrc wave=sine freq=%s "

412 "! audioconvert "

413 "! audioresample "

414 "! autoaudiosink" % int(frequency)

415 )

416 self.logger = logger

417 if self.logger: 417 ↛ 419 line 417 didn't jump to line 419 because the condition on line 417 was always true

418 self.logger.debug(self.pipeline_description)

419 self.pipeline = Gst.parse_launch(self.pipeline_description)

420

421

422 class Recorder(GstAudioObject):

423 def __init__(

424 self,

425 output_file,

426 bins=BINS,

427 sampling_frequency=SAMPLING_FREQUENCY,

428 fft_interval=FFT_INTERVAL,

429 logger=None,

430 ):

431 super(Recorder, self).__init__()

432 pipeline_description = """autoaudiosrc

433 ! queue

434 ! level message=true

435 ! audioconvert

436 ! audio/x-raw, channels=1, rate=(int)%(rate)s

437 ! audioresample

438 ! spectrum interval=%(fft_interval)s bands = %(bands)s

439 ! wavenc

440 ! filesink location=%(file)s""" % {

441 "bands": bins,

442 "rate": sampling_frequency,

443 "fft_interval": fft_interval,

444 "file": output_file,

445 }

446 self.logger = logger

447 if self.logger: 447 ↛ 449 line 447 didn't jump to line 449 because the condition on line 447 was always true

448 self.logger.debug(pipeline_description)

449 self.pipeline = Gst.parse_launch(pipeline_description)

450

451 def register_message_handler(self, handler_method):

452 if self.logger: 452 ↛ 455 line 452 didn't jump to line 455 because the condition on line 452 was always true

453 message = "Registering message handler: %s" % handler_method

454 self.logger.debug(message)

455 self.bus = self.pipeline.get_bus()

456 self.bus.add_signal_watch()

457 self.bus.connect("message", handler_method)

458

459

460 def parse_spectrum_message_structure(struct_string):

461 # First let's jsonize this

462 # This is the message name, which we don't need

463 text = struct_string.replace("spectrum, ", "")

464 # name/value separator in json is : and not =

465 text = text.replace("=", ": ")

466 # Mutate the {} array notation from the structure to

467 # [] notation for json.

468 text = text.replace("{", "[")

469 text = text.replace("}", "]")

470 # Remove a few stray semicolons that aren't needed

471 text = text.replace(";", "")

472 # Remove the data type fields, as json doesn't need them

473 text = re.sub(r"\(.+?\)", "", text)

474 # double-quote the identifiers

475 text = re.sub(r"([\w-]+):", r'"\1":', text)

476 # Wrap the whole thing in brackets

477 text = "{" + text + "}"

478 # Try to parse and return something sensible here, even if

479 # the data was unparsable.

480 try:

481 return json.loads(text)

482 except ValueError as e:

483 logging.error(repr(e))

484 return None

485

486

487 def process_arguments(args=sys.argv[1:]):

488 description = """

489 Plays a single frequency through the default output, then records on

490 the default input device. Analyzes the recorded signal to test for

491 presence of the played frequency, if present it exits with success.

492 """

493 parser = argparse.ArgumentParser(description=description)

494 parser.add_argument(

495 "-t",

496 "--time",

497 dest="test_duration",

498 action="store",

499 default=30,

500 type=int,

501 help="""Maximum test duration, default %(default)s seconds.

502 It may exit sooner if it determines it has enough data.""",

503 )

504 parser.add_argument(

505 "-a",

506 "--audio",

507 action="store",

508 default="/dev/null",

509 type=str,

510 help="File to save recorded audio in .wav format",

511 )

512 parser.add_argument(

513 "-q",

514 "--quiet",

515 action="store_true",

516 default=False,

517 help="Be quiet, no output unless there's an error.",

518 )

519 parser.add_argument(

520 "-d",

521 "--debug",

522 action="store_true",

523 default=False,

524 help="Debugging output",

525 )

526 parser.add_argument(

527 "-f",

528 "--frequency",

529 action="store",

530 default=DEFAULT_TEST_FREQUENCY,

531 type=int,

532 help="Frequency for test signal, default %(default)s Hz",

533 )

534 parser.add_argument(

535 "-u",

536 "--spectrum",

537 action="store",

538 type=str,

539 help="""File to save spectrum information for plotting

540 (one frequency/magnitude pair per line)""",

541 )

542 return parser.parse_args(args)

543

544

545 def run_test(args):

546

547 # Setup logging

548 level = logging.INFO

549 if args.debug:

550 level = logging.DEBUG

551 if args.quiet:

552 level = logging.ERROR

553 logging.basicConfig(level=level)

554 try:

555 # Launches recording pipeline. I need to hook up into the gst

556 # messages.

557 recorder = Recorder(output_file=args.audio, logger=logging)

558 # Just launches the playing pipeline

559 player = Player(frequency=args.frequency, logger=logging)

560 except GObject.GError as excp:

561 logging.critical("Unable to initialize GStreamer pipelines: %s", excp)

562 sys.exit(127)

563

564 # This just receives a process feedback and tells me how much to change to

565 # achieve the setpoint

566 pidctrl = PIDController(

567 kp=0.7, ki=0.01, kd=0.01, setpoint=REC_LEVEL_RANGE[0]

568 )

569 pidctrl.set_change_limit(5)

570 # This gathers spectrum data.

571 analyzer = SpectrumAnalyzer(

572 points=BINS, sampling_frequency=SAMPLING_FREQUENCY

573 )

574

575 recorder.volumecontroller = VolumeController(type="input", logger=logging)

576 recorder.volumecontroller.set_volume(0)

577 recorder.volumecontroller.mute(False)

578

579 player.volumecontroller = VolumeController(type="output", logger=logging)

580 player.volumecontroller.set_volume(PLAY_VOLUME)

581 player.volumecontroller.mute(False)

582

583 # This handles the messages from gstreamer and orchestrates

584 # the passed volume controllers, pid controller and spectrum analyzer

585 # accordingly.

586 gmh = GStreamerMessageHandler(

587 rec_level_range=REC_LEVEL_RANGE,

588 logger=logging,

589 volumecontroller=recorder.volumecontroller,

590 pidcontroller=pidctrl,

591 spectrum_analyzer=analyzer,

592 )

593

594 # I need to tell the recorder which method will handle messages.

595 recorder.register_message_handler(gmh.bus_message_handler)

596

597 # Create the loop and add a few triggers

598 loop = GLib.MainLoop()

599 GLib.timeout_add_seconds(0, player.start)

600 GLib.timeout_add_seconds(0, recorder.start)

601 GLib.timeout_add_seconds(args.test_duration, loop.quit)

602

603 # Tell the gmh which method to call when enough samples are collected

604 gmh.set_quit_method(loop.quit)

605

606 loop.run()

607

608 # When the loop ends, set things back to reasonable states

609 player.stop()

610 recorder.stop()

611 player.volumecontroller.set_volume(50)

612 recorder.volumecontroller.set_volume(10)

613

614 # See if data gathering was successful.

615 test_band = analyzer.frequency_band_for(args.frequency)

616 candidate_bands = analyzer.frequencies_with_peak_magnitude(

617 MAGNITUDE_THRESHOLD

618 )

619 for band in candidate_bands:

620 logging.debug(

621 "Band (%.2f,%.2f) contains a magnitude peak"

622 % analyzer.frequencies_for_band(band)

623 )

624 if test_band in candidate_bands:

625 freqs_for_band = analyzer.frequencies_for_band(test_band)

626 logging.info(

627 "PASS: Test frequency of %s in band (%.2f, %.2f) "

628 "which contains a magnitude peak"

629 % ((args.frequency,) + freqs_for_band)

630 )

631 return_value = 0

632 else:

633 logging.info(

634 "FAIL: Test frequency of %s is not in one of the "

635 "bands with magnitude peaks" % args.frequency

636 )

637 return_value = 1

638 # Is the microphone broken?

639 if len(set(analyzer.spectrum)) <= 1: 639 ↛ 645 line 639 didn't jump to line 645 because the condition on line 639 was always true

640 logging.info(

641 "WARNING: Microphone seems broken, didn't even "

642 "record ambient noise"

643 )

644

645 if args.spectrum:

646 logging.info("Saving spectrum data for plotting as %s" % args.spectrum)

647 if not FileDumper().write_to_file( 647 ↛ 654 line 647 didn't jump to line 654 because the condition on line 647 was never true

648 args.spectrum,

649 [

650 "%s,%s" % t

651 for t in zip(analyzer.frequencies, analyzer.spectrum)

652 ],

653 ):

654 logging.error(

655 "Couldn't save spectrum data for plotting", file=sys.stderr

656 )

657

658 return return_value

659

660

661 def main():

662 # Get arguments.

663 args = process_arguments()

664

665 return run_test(args)

666

667

668 if __name__ == "__main__":

669 sys.exit(main())