bazarr/libs/auditok/workers.py

import os
import sys
from tempfile import NamedTemporaryFile
from abc import ABCMeta, abstractmethod
from threading import Thread
from datetime import datetime, timedelta
from collections import namedtuple
import wave
import subprocess
from queue import Queue, Empty
from .io import _guess_audio_format
from .util import AudioDataSource, make_duration_formatter
from .core import split
from .exceptions import (
    EndOfProcessing,
    AudioEncodingError,
    AudioEncodingWarning,
)


_STOP_PROCESSING = "STOP_PROCESSING"
_Detection = namedtuple("_Detection", "id start end duration")


def _run_subprocess(command):
    try:
        with subprocess.Popen(
            command,
            stdin=open(os.devnull, "rb"),
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        ) as proc:
            stdout, stderr = proc.communicate()
            return proc.returncode, stdout, stderr
    except Exception:
        err_msg = "Couldn't export audio using command: '{}'".format(command)
        raise AudioEncodingError(err_msg)


class Worker(Thread, metaclass=ABCMeta):
    def __init__(self, timeout=0.5, logger=None):
        self._timeout = timeout
        self._logger = logger
        self._inbox = Queue()
        Thread.__init__(self)

    def run(self):
        while True:
            message = self._get_message()
            if message == _STOP_PROCESSING:
                break
            if message is not None:
                self._process_message(message)
        self._post_process()

    @abstractmethod
    def _process_message(self, message):
        """Process incoming messages"""

    def _post_process(self):
        pass

    def _log(self, message):
        self._logger.info(message)

    def _stop_requested(self):
        try:
            message = self._inbox.get_nowait()
            if message == _STOP_PROCESSING:
                return True
        except Empty:
            return False

    def stop(self):
        self.send(_STOP_PROCESSING)
        self.join()

    def send(self, message):
        self._inbox.put(message)

    def _get_message(self):
        try:
            message = self._inbox.get(timeout=self._timeout)
            return message
        except Empty:
            return None


class TokenizerWorker(Worker, AudioDataSource):
    def __init__(self, reader, observers=None, logger=None, **kwargs):
        self._observers = observers if observers is not None else []
        self._reader = reader
        self._audio_region_gen = split(self, **kwargs)
        self._detections = []
        self._log_format = "[DET]: Detection {0.id} (start: {0.start:.3f}, "
        self._log_format += "end: {0.end:.3f}, duration: {0.duration:.3f})"
        Worker.__init__(self, timeout=0.2, logger=logger)

    def _process_message(self):
        pass

    @property
    def detections(self):
        return self._detections

    def _notify_observers(self, message):
        for observer in self._observers:
            observer.send(message)

    def run(self):
        self._reader.open()
        start_processing_timestamp = datetime.now()
        for _id, audio_region in enumerate(self._audio_region_gen, start=1):
            timestamp = start_processing_timestamp + timedelta(
                seconds=audio_region.meta.start
            )
            audio_region.meta.timestamp = timestamp
            detection = _Detection(
                _id,
                audio_region.meta.start,
                audio_region.meta.end,
                audio_region.duration,
            )
            self._detections.append(detection)
            if self._logger is not None:
                message = self._log_format.format(detection)
                self._log(message)
            self._notify_observers((_id, audio_region))
        self._notify_observers(_STOP_PROCESSING)
        self._reader.close()

    def start_all(self):
        for observer in self._observers:
            observer.start()
        self.start()

    def stop_all(self):
        self.stop()
        for observer in self._observers:
            observer.stop()
        self._reader.close()

    def read(self):
        if self._stop_requested():
            return None
        else:
            return self._reader.read()

    def __getattr__(self, name):
        return getattr(self._reader, name)


class StreamSaverWorker(Worker):
    def __init__(
        self,
        audio_reader,
        filename,
        export_format=None,
        cache_size_sec=0.5,
        timeout=0.2,
    ):
        self._reader = audio_reader
        sample_size_bytes = self._reader.sw * self._reader.ch
        self._cache_size = cache_size_sec * self._reader.sr * sample_size_bytes
        self._output_filename = filename
        self._export_format = _guess_audio_format(export_format, filename)
        if self._export_format is None:
            self._export_format = "wav"
        self._init_output_stream()
        self._exported = False
        self._cache = []
        self._total_cached = 0
        Worker.__init__(self, timeout=timeout)

    def _get_non_existent_filename(self):
        filename = self._output_filename + ".wav"
        i = 0
        while os.path.exists(filename):
            i += 1
            filename = self._output_filename + "({}).wav".format(i)
        return filename

    def _init_output_stream(self):
        if self._export_format != "wav":
            self._tmp_output_filename = self._get_non_existent_filename()
        else:
            self._tmp_output_filename = self._output_filename
        self._wfp = wave.open(self._tmp_output_filename, "wb")
        self._wfp.setframerate(self._reader.sr)
        self._wfp.setsampwidth(self._reader.sw)
        self._wfp.setnchannels(self._reader.ch)

    @property
    def sr(self):
        return self._reader.sampling_rate

    @property
    def sw(self):
        return self._reader.sample_width

    @property
    def ch(self):
        return self._reader.channels

    def __del__(self):
        self._post_process()

        if (
            (self._tmp_output_filename != self._output_filename)
            and self._exported
            and os.path.exists(self._tmp_output_filename)
        ):
            os.remove(self._tmp_output_filename)

    def _process_message(self, data):
        self._cache.append(data)
        self._total_cached += len(data)
        if self._total_cached >= self._cache_size:
            self._write_cached_data()

    def _post_process(self):
        while True:
            try:
                data = self._inbox.get_nowait()
                if data != _STOP_PROCESSING:
                    self._cache.append(data)
                    self._total_cached += len(data)
            except Empty:
                break
        self._write_cached_data()
        self._wfp.close()

    def _write_cached_data(self):
        if self._cache:
            data = b"".join(self._cache)
            self._wfp.writeframes(data)
            self._cache = []
            self._total_cached = 0

    def open(self):
        self._reader.open()

    def close(self):
        self._reader.close()
        self.stop()

    def rewind(self):
        # ensure compatibility with AudioDataSource with record=True
        pass

    @property
    def data(self):
        with wave.open(self._tmp_output_filename, "rb") as wfp:
            return wfp.readframes(-1)

    def save_stream(self):
        if self._exported:
            return self._output_filename

        if self._export_format in ("raw", "wav"):
            if self._export_format == "raw":
                self._export_raw()
            self._exported = True
            return self._output_filename
        try:
            self._export_with_ffmpeg_or_avconv()
        except AudioEncodingError:
            try:
                self._export_with_sox()
            except AudioEncodingError:
                warn_msg = "Couldn't save audio data in the desired format "
                warn_msg += "'{}'. Either none of 'ffmpeg', 'avconv' or 'sox' "
                warn_msg += "is installed or this format is not recognized.\n"
                warn_msg += "Audio file was saved as '{}'"
                raise AudioEncodingWarning(
                    warn_msg.format(
                        self._export_format, self._tmp_output_filename
                    )
                )
        finally:
            self._exported = True
        return self._output_filename

    def _export_raw(self):
        with open(self._output_filename, "wb") as wfp:
            wfp.write(self.data)

    def _export_with_ffmpeg_or_avconv(self):
        command = [
            "-y",
            "-f",
            "wav",
            "-i",
            self._tmp_output_filename,
            "-f",
            self._export_format,
            self._output_filename,
        ]
        returncode, stdout, stderr = _run_subprocess(["ffmpeg"] + command)
        if returncode != 0:
            returncode, stdout, stderr = _run_subprocess(["avconv"] + command)
            if returncode != 0:
                raise AudioEncodingError(stderr)
        return stdout, stderr

    def _export_with_sox(self):
        command = [
            "sox",
            "-t",
            "wav",
            self._tmp_output_filename,
            self._output_filename,
        ]
        returncode, stdout, stderr = _run_subprocess(command)
        if returncode != 0:
            raise AudioEncodingError(stderr)
        return stdout, stderr

    def close_output(self):
        self._wfp.close()

    def read(self):
        data = self._reader.read()
        if data is not None:
            self.send(data)
        else:
            self.send(_STOP_PROCESSING)
        return data

    def __getattr__(self, name):
        if name == "data":
            return self.data
        return getattr(self._reader, name)


class PlayerWorker(Worker):
    def __init__(self, player, progress_bar=False, timeout=0.2, logger=None):
        self._player = player
        self._progress_bar = progress_bar
        self._log_format = "[PLAY]: Detection {id} played"
        Worker.__init__(self, timeout=timeout, logger=logger)

    def _process_message(self, message):
        _id, audio_region = message
        if self._logger is not None:
            message = self._log_format.format(id=_id)
            self._log(message)
        audio_region.play(
            player=self._player, progress_bar=self._progress_bar, leave=False
        )


class RegionSaverWorker(Worker):
    def __init__(
        self,
        filename_format,
        audio_format=None,
        timeout=0.2,
        logger=None,
        **audio_parameters
    ):
        self._filename_format = filename_format
        self._audio_format = audio_format
        self._audio_parameters = audio_parameters
        self._debug_format = "[SAVE]: Detection {id} saved as '{filename}'"
        Worker.__init__(self, timeout=timeout, logger=logger)

    def _process_message(self, message):
        _id, audio_region = message
        filename = self._filename_format.format(
            id=_id,
            start=audio_region.meta.start,
            end=audio_region.meta.end,
            duration=audio_region.duration,
        )
        filename = audio_region.save(
            filename, self._audio_format, **self._audio_parameters
        )
        if self._logger:
            message = self._debug_format.format(id=_id, filename=filename)
            self._log(message)


class CommandLineWorker(Worker):
    def __init__(self, command, timeout=0.2, logger=None):
        self._command = command
        Worker.__init__(self, timeout=timeout, logger=logger)
        self._debug_format = "[COMMAND]: Detection {id} command: '{command}'"

    def _process_message(self, message):
        _id, audio_region = message
        with NamedTemporaryFile(delete=False) as file:
            filename = audio_region.save(file.name, audio_format="wav")
            command = self._command.format(file=filename)
            os.system(command)
            if self._logger is not None:
                message = self._debug_format.format(id=_id, command=command)
                self._log(message)


class PrintWorker(Worker):
    def __init__(
        self,
        print_format="{start} {end}",
        time_format="%S",
        timestamp_format="%Y/%m/%d %H:%M:%S.%f",
        timeout=0.2,
    ):

        self._print_format = print_format
        self._format_time = make_duration_formatter(time_format)
        self._timestamp_format = timestamp_format
        self.detections = []
        Worker.__init__(self, timeout=timeout)

    def _process_message(self, message):
        _id, audio_region = message
        timestamp = audio_region.meta.timestamp
        timestamp = timestamp.strftime(self._timestamp_format)
        text = self._print_format.format(
            id=_id,
            start=self._format_time(audio_region.meta.start),
            end=self._format_time(audio_region.meta.end),
            duration=self._format_time(audio_region.duration),
            timestamp=timestamp,
        )
        print(text)