library/enviroplus/noise.py

   1 import sounddevice
   2 import numpy
   3 import math
   4
   5 class Noise():
   6     def __init__(
   7         self,
   8         sample_rate=16000,
   9         duration=0.5):
  10
  11         self.duration = duration
  12         self.sample_rate = sample_rate
  13
  14     def get_amplitudes_at_frequency_ranges(self, ranges):
  15         recording = self._record()
  16         magnitude = numpy.abs(numpy.fft.rfft(recording[:, 0], n=self.sample_rate))
  17         result = []
  18         for r in ranges:
  19             start, end = r
  20             result.append(numpy.mean(magnitude[start:end]))
  21         return result
  22
  23     def get_amplitude_at_frequency_range(self, start, end):
  24         n = self.sample_rate // 2
  25         if start > n or end > n:
  26             raise ValueError("Maxmimum frequency is {}".format(n))
  27
  28         recording = self._record()
  29         magnitude = numpy.abs(numpy.fft.rfft(recording[:, 0], n=self.sample_rate))
  30         return numpy.mean(magnitude[start:end])
  31
  32     def get_noise_profile(
  33         self,
  34         noise_floor=100,
  35         low=0.12,
  36         mid=0.36,
  37         high=None):
  38
  39         if high is None:
  40             high = 1.0 - low - mid
  41
  42         recording = self._record()
  43         magnitude = numpy.abs(numpy.fft.rfft(recording[:, 0], n=self.sample_rate))
  44
  45         sample_count = (self.sample_rate // 2) - noise_floor
  46
  47         mid_start = noise_floor + int(sample_count * low)
  48         high_start = mid_start + int(sample_count * mid)
  49         noise_ceiling = high_start + int(sample_count * high)
  50
  51         amp_low = numpy.mean(magnitude[self.noise_floor:mid_start])
  52         amp_mid = numpy.mean(magnitude[mid_start:high_start])
  53         amp_high = numpy.mean(magnitude[high_start:noise_ceiling])
  54         amp_total = (low + mid + high) / 3.0
  55
  56         return amp_low, amp_mid, amp_high, amp_total
  57
  58     def _record(self):
  59         return sounddevice.rec(
  60             int(self.duration * self.sample_rate),
  61             samplerate=self.sample_rate,
  62             blocking=True,
  63             channels=1,
  64             dtype='float64'
  65         )
  66