.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "tutorials/device_avsr.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note Click :ref:`here ` to download the full example code .. rst-class:: sphx-glr-example-title .. _sphx_glr_tutorials_device_avsr.py: Device AV-ASR with Emformer RNN-T ================================= **Author**: `Pingchuan Ma `__, `Moto Hira `__. This tutorial shows how to run on-device audio-visual speech recognition (AV-ASR, or AVSR) with TorchAudio on a streaming device input, i.e. microphone on laptop. AV-ASR is the task of transcribing text from audio and visual streams, which has recently attracted a lot of research attention due to its robustness against noise. .. note:: This tutorial requires ffmpeg, sentencepiece, mediapipe, opencv-python and scikit-image libraries. There are multiple ways to install ffmpeg libraries. If you are using Anaconda Python distribution, ``conda install -c conda-forge 'ffmpeg<7'`` will install compatible FFmpeg libraries. You can run ``pip install sentencepiece mediapipe opencv-python scikit-image`` to install the other libraries mentioned. .. note:: To run this tutorial, please make sure you are in the `tutorial` folder. .. note:: We tested the tutorial on torchaudio version 2.0.2 on Macbook Pro (M1 Pro). .. GENERATED FROM PYTHON SOURCE LINES 37-45 .. code-block:: default import numpy as np import sentencepiece as spm import torch import torchaudio import torchvision .. GENERATED FROM PYTHON SOURCE LINES 46-59 Overview -------- The real-time AV-ASR system is presented as follows, which consists of three components, a data collection module, a pre-processing module and an end-to-end model. The data collection module is hardware, such as a microphone and camera. Its role is to collect information from the real world. Once the information is collected, the pre-processing module location and crop out face. Next, we feed the raw audio stream and the pre-processed video stream into our end-to-end model for inference. .. image:: https://download.pytorch.org/torchaudio/doc-assets/avsr/overview.png .. GENERATED FROM PYTHON SOURCE LINES 62-72 1. Data acquisition ------------------- Firstly, we define the function to collect videos from microphone and camera. To be specific, we use :py:class:`~torchaudio.io.StreamReader` class for the purpose of data collection, which supports capturing audio/video from microphone and camera. For the detailed usage of this class, please refer to the `tutorial <./streamreader_basic_tutorial.html>`__. .. GENERATED FROM PYTHON SOURCE LINES 72-114 .. code-block:: default def stream(q, format, option, src, segment_length, sample_rate): print("Building StreamReader...") streamer = torchaudio.io.StreamReader(src=src, format=format, option=option) streamer.add_basic_video_stream(frames_per_chunk=segment_length, buffer_chunk_size=500, width=600, height=340) streamer.add_basic_audio_stream(frames_per_chunk=segment_length * 640, sample_rate=sample_rate) print(streamer.get_src_stream_info(0)) print(streamer.get_src_stream_info(1)) print("Streaming...") print() for (chunk_v, chunk_a) in streamer.stream(timeout=-1, backoff=1.0): q.put([chunk_v, chunk_a]) class ContextCacher: def __init__(self, segment_length: int, context_length: int, rate_ratio: int): self.segment_length = segment_length self.context_length = context_length self.context_length_v = context_length self.context_length_a = context_length * rate_ratio self.context_v = torch.zeros([self.context_length_v, 3, 340, 600]) self.context_a = torch.zeros([self.context_length_a, 1]) def __call__(self, chunk_v, chunk_a): if chunk_v.size(0) < self.segment_length: chunk_v = torch.nn.functional.pad(chunk_v, (0, 0, 0, 0, 0, 0, 0, self.segment_length - chunk_v.size(0))) if chunk_a.size(0) < self.segment_length * 640: chunk_a = torch.nn.functional.pad(chunk_a, (0, 0, 0, self.segment_length * 640 - chunk_a.size(0))) if self.context_length == 0: return chunk_v.float(), chunk_a.float() else: chunk_with_context_v = torch.cat((self.context_v, chunk_v)) chunk_with_context_a = torch.cat((self.context_a, chunk_a)) self.context_v = chunk_v[-self.context_length_v :] self.context_a = chunk_a[-self.context_length_a :] return chunk_with_context_v.float(), chunk_with_context_a.float() .. GENERATED FROM PYTHON SOURCE LINES 115-140 2. Pre-processing ----------------- Before feeding the raw stream into our model, each video sequence has to undergo a specific pre-processing procedure. This involves three critical steps. The first step is to perform face detection. Following that, each individual frame is aligned to a referenced frame, commonly known as the mean face, in order to normalize rotation and size differences across frames. The final step in the pre-processing module is to crop the face region from the aligned face image. .. list-table:: :widths: 25 25 25 25 :header-rows: 0 * - .. image:: https://download.pytorch.org/torchaudio/doc-assets/avsr/original.gif - .. image:: https://download.pytorch.org/torchaudio/doc-assets/avsr/detected.gif - .. image:: https://download.pytorch.org/torchaudio/doc-assets/avsr/transformed.gif - .. image:: https://download.pytorch.org/torchaudio/doc-assets/avsr/cropped.gif * - 0. Original - 1. Detected - 2. Transformed - 3. Cropped .. GENERATED FROM PYTHON SOURCE LINES 140-183 .. code-block:: default import sys sys.path.insert(0, "../../examples") from avsr.data_prep.detectors.mediapipe.detector import LandmarksDetector from avsr.data_prep.detectors.mediapipe.video_process import VideoProcess class FunctionalModule(torch.nn.Module): def __init__(self, functional): super().__init__() self.functional = functional def forward(self, input): return self.functional(input) class Preprocessing(torch.nn.Module): def __init__(self): super().__init__() self.landmarks_detector = LandmarksDetector() self.video_process = VideoProcess() self.video_transform = torch.nn.Sequential( FunctionalModule( lambda n: [(lambda x: torchvision.transforms.functional.resize(x, 44, antialias=True))(i) for i in n] ), FunctionalModule(lambda x: torch.stack(x)), torchvision.transforms.Normalize(0.0, 255.0), torchvision.transforms.Grayscale(), torchvision.transforms.Normalize(0.421, 0.165), ) def forward(self, audio, video): video = video.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8) landmarks = self.landmarks_detector(video) video = self.video_process(video, landmarks) video = torch.tensor(video).permute(0, 3, 1, 2).float() video = self.video_transform(video) audio = audio.mean(axis=-1, keepdim=True) return audio, video .. GENERATED FROM PYTHON SOURCE LINES 184-198 3. Building inference pipeline ------------------------------ The next step is to create components required for pipeline. We use convolutional-based front-ends to extract features from both the raw audio and video streams. These features are then passed through a two-layer MLP for fusion. For our transducer model, we leverage the TorchAudio library, which incorporates an encoder (Emformer), a predictor, and a joint network. The architecture of the proposed AV-ASR model is illustrated as follows. .. image:: https://download.pytorch.org/torchaudio/doc-assets/avsr/architecture.png .. GENERATED FROM PYTHON SOURCE LINES 198-258 .. code-block:: default class SentencePieceTokenProcessor: def __init__(self, sp_model): self.sp_model = sp_model self.post_process_remove_list = { self.sp_model.unk_id(), self.sp_model.eos_id(), self.sp_model.pad_id(), } def __call__(self, tokens, lstrip: bool = True) -> str: filtered_hypo_tokens = [ token_index for token_index in tokens[1:] if token_index not in self.post_process_remove_list ] output_string = "".join(self.sp_model.id_to_piece(filtered_hypo_tokens)).replace("\u2581", " ") if lstrip: return output_string.lstrip() else: return output_string class InferencePipeline(torch.nn.Module): def __init__(self, preprocessor, model, decoder, token_processor): super().__init__() self.preprocessor = preprocessor self.model = model self.decoder = decoder self.token_processor = token_processor self.state = None self.hypotheses = None def forward(self, audio, video): audio, video = self.preprocessor(audio, video) feats = self.model(audio.unsqueeze(0), video.unsqueeze(0)) length = torch.tensor([feats.size(1)], device=audio.device) self.hypotheses, self.state = self.decoder.infer(feats, length, 10, state=self.state, hypothesis=self.hypotheses) transcript = self.token_processor(self.hypotheses[0][0], lstrip=False) return transcript def _get_inference_pipeline(model_path, spm_model_path): model = torch.jit.load(model_path) model.eval() sp_model = spm.SentencePieceProcessor(model_file=spm_model_path) token_processor = SentencePieceTokenProcessor(sp_model) decoder = torchaudio.models.RNNTBeamSearch(model.model, sp_model.get_piece_size()) return InferencePipeline( preprocessor=Preprocessing(), model=model, decoder=decoder, token_processor=token_processor, ) .. GENERATED FROM PYTHON SOURCE LINES 259-269 4. The main process ------------------- The execution flow of the main process is as follows: 1. Initialize the inference pipeline. 2. Launch data acquisition subprocess. 3. Run inference. 4. Clean up .. GENERATED FROM PYTHON SOURCE LINES 269-328 .. code-block:: default from torchaudio.utils import download_asset def main(device, src, option=None): print("Building pipeline...") model_path = download_asset("tutorial-assets/device_avsr_model.pt") spm_model_path = download_asset("tutorial-assets/spm_unigram_1023.model") pipeline = _get_inference_pipeline(model_path, spm_model_path) BUFFER_SIZE = 32 segment_length = 8 context_length = 4 sample_rate = 19200 frame_rate = 30 rate_ratio = sample_rate // frame_rate cacher = ContextCacher(BUFFER_SIZE, context_length, rate_ratio) import torch.multiprocessing as mp ctx = mp.get_context("spawn") @torch.inference_mode() def infer(): num_video_frames = 0 video_chunks = [] audio_chunks = [] while True: chunk_v, chunk_a = q.get() num_video_frames += chunk_a.size(0) // 640 video_chunks.append(chunk_v) audio_chunks.append(chunk_a) if num_video_frames < BUFFER_SIZE: continue video = torch.cat(video_chunks) audio = torch.cat(audio_chunks) video, audio = cacher(video, audio) pipeline.state, pipeline.hypotheses = None, None transcript = pipeline(audio, video.float()) print(transcript, end="", flush=True) num_video_frames = 0 video_chunks = [] audio_chunks = [] q = ctx.Queue() p = ctx.Process(target=stream, args=(q, device, option, src, segment_length, sample_rate)) p.start() infer() p.join() if __name__ == "__main__": main( device="avfoundation", src="0:1", option={"framerate": "30", "pixel_format": "rgb24"}, ) .. GENERATED FROM PYTHON SOURCE LINES 329-339 .. code:: Building pipeline... Building StreamReader... SourceVideoStream(media_type='video', codec='rawvideo', codec_long_name='raw video', format='uyvy422', bit_rate=0, num_frames=0, bits_per_sample=0, metadata={}, width=1552, height=1552, frame_rate=1000000.0) SourceAudioStream(media_type='audio', codec='pcm_f32le', codec_long_name='PCM 32-bit floating point little-endian', format='flt', bit_rate=1536000, num_frames=0, bits_per_sample=0, metadata={}, sample_rate=48000.0, num_channels=1) Streaming... hello world .. GENERATED FROM PYTHON SOURCE LINES 342-344 Tag: :obj:`torchaudio.io` .. rst-class:: sphx-glr-timing **Total running time of the script:** ( 0 minutes 0.000 seconds) .. _sphx_glr_download_tutorials_device_avsr.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: device_avsr.py ` .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: device_avsr.ipynb ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_