\n", + "\n" + ], + "text/plain": [ + "\n" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n", + " | Name | Type | Params | In sizes | Out sizes \n", + "----------------------------------------------------------------------------------------------------------------------\n", + "0 | mfcc | MFCC | 0 | [32, 1, 32000] | [32, 1, 40, 161] \n", + "1 | lstm | LSTM | 18.9 K | [32, 161, 40] | [[32, 161, 64], [[2, 32, 32], [2, 32, 32]]]\n", + "2 | classifier | Linear | 195 | [32, 161, 64] | [32, 161, 3] \n", + "3 | activation | Sigmoid | 0 | [32, 161, 3] | [32, 161, 3] \n", + "4 | validation_metric | MetricCollection | 0 | ? | ? \n", + "----------------------------------------------------------------------------------------------------------------------\n", + "19.1 K Trainable params\n", + "0 Non-trainable params\n", + "19.1 K Total params\n", + "0.077 Total estimated model params size (MB)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + " - 27.32% of all chunks contain no speech at all.\n", + " - 66.79% contain 1 speaker or less\n", + " - 91.25% contain 2 speakers or less\n", + " - 98.04% contain 3 speakers or less\n", + " - 100.00% contain 4 speakers or less\n", + "Setting `max_speakers_per_chunk` to 3. You can override this value (or avoid this estimation step) by passing `max_speakers_per_chunk=3` to the task constructor.\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "423b6758f1c84f3b8ece70744862fb66", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Sanity Checking: 0it [00:00, ?it/s]" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "e85cb23d204c4876a82a8202edc273c5", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Training: 0it [00:00, ?it/s]" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "b8f5929201654acf83fbca919a5ac63b", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Validation: 0it [00:00, ?it/s]" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "`Trainer.fit` stopped: `max_epochs=1` reached.\n" + ] + } + ], "source": [ - "from pyannote.audio.tasks import Segmentation\n", - "seg = Segmentation(protocol, duration=2., batch_size=32, num_workers=4)\n", + "from pyannote.audio.tasks import SpeakerDiarization\n", + "seg = SpeakerDiarization(protocol, duration=2., batch_size=32, num_workers=4)\n", "model = SimpleSegmentationModel(task=seg)\n", "trainer = pl.Trainer(max_epochs=1)\n", "_ = trainer.fit(model)" ] }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Speaker tracking" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from pyannote.audio.tasks import SpeakerTracking\n", - "spk = SpeakerTracking(protocol, duration=2., batch_size=32, num_workers=4)\n", - "model = SimpleSegmentationModel(task=spk)\n", - "trainer = pl.Trainer(max_epochs=1)\n", - "_ = trainer.fit(model)" - ] - }, { "cell_type": "markdown", "metadata": {}, @@ -141,9 +420,68 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 7, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "GPU available: True (mps), used: False\n", + "TPU available: False, using: 0 TPU cores\n", + "IPU available: False, using: 0 IPUs\n", + "HPU available: False, using: 0 HPUs\n", + "Loading Debug.SpeakerDiarization.Debug training labels: 10file [00:00, 363.55file/s]\n", + "\n", + " | Name | Type | Params | In sizes | Out sizes \n", + "----------------------------------------------------------------------------------------------------------------------\n", + "0 | mfcc | MFCC | 0 | [32, 1, 32000] | [32, 1, 40, 161] \n", + "1 | lstm | LSTM | 18.9 K | [32, 161, 40] | [[32, 161, 64], [[2, 32, 32], [2, 32, 32]]]\n", + "2 | loss_func | ArcFaceLoss | 832 | ? | ? \n", + "3 | validation_metric | MetricCollection | 0 | ? | ? \n", + "----------------------------------------------------------------------------------------------------------------------\n", + "19.8 K Trainable params\n", + "0 Non-trainable params\n", + "19.8 K Total params\n", + "0.079 Total estimated model params size (MB)\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "b34c5cbe6c9c409fb5572b0be9598076", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Sanity Checking: 0it [00:00, ?it/s]" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "71077814aa154ba3b4892e8be13a98fe", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Training: 0it [00:00, ?it/s]" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "`Trainer.fit` stopped: `max_epochs=1` reached.\n" + ] + } + ], "source": [ "from pyannote.audio.tasks import SpeakerEmbedding\n", "emb = SpeakerEmbedding(protocol, duration=2., num_workers=4)\n", @@ -169,7 +507,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.5" + "version": "3.9.13" } }, "nbformat": 4, diff --git a/notebook/inference.ipynb b/notebook/inference.ipynb index 77e02a27c..b33e64282 100644 --- a/notebook/inference.ipynb +++ b/notebook/inference.ipynb @@ -104,19 +104,6 @@ "scores = inference.crop(dev_file, Segment(10, 15))" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# inference with progress bar\n", - "inference = Inference(model, step=0.1, batch_size=32, progress_hook='Processing...')\n", - "scores = inference(dev_file)\n", - "inference = Inference(model, step=0.1, batch_size=1, progress_hook=True)\n", - "scores = inference.crop(dev_file, Segment(10, 15))" - ] - }, { "cell_type": "markdown", "metadata": {}, diff --git a/pyannote/audio/cli/evaluate.py b/pyannote/audio/cli/evaluate.py index 57f5c0149..a5ab682c5 100644 --- a/pyannote/audio/cli/evaluate.py +++ b/pyannote/audio/cli/evaluate.py @@ -53,14 +53,17 @@ def evaluate(cfg: DictConfig) -> Optional[float]: main_task = progress.add_task(protocol.name, total=len(files)) file_task = progress.add_task("Processing", total=1.0) - def progress_hook(completed: int, total: int): + def progress_hook(completed: int = None, total: int = None): progress.update(file_task, completed=completed / total) - inference = Inference(model, device=device, progress_hook=progress_hook) + inference = Inference(model, device=device) warm_up = cfg.warm_up / inference.duration def hypothesis(file: ProtocolFile): - return Inference.trim(binarize(inference(file)), warm_up=(warm_up, warm_up)) + return Inference.trim( + binarize(inference(file, hook=progress_hook)), + warm_up=(warm_up, warm_up), + ) metric = DiscreteDiarizationErrorRate() diff --git a/pyannote/audio/cli/pretrained.py b/pyannote/audio/cli/pretrained.py index 95a439b3f..8b8aae587 100644 --- a/pyannote/audio/cli/pretrained.py +++ b/pyannote/audio/cli/pretrained.py @@ -26,6 +26,4 @@ def pretrained(checkpoint: Text): - return Model.from_pretrained( - checkpoint, map_location=lambda storage, loc: storage - ) + return Model.from_pretrained(checkpoint, map_location=lambda storage, loc: storage) diff --git a/pyannote/audio/cli/train.py b/pyannote/audio/cli/train.py index 6e04fcde5..9ab8b1658 100644 --- a/pyannote/audio/cli/train.py +++ b/pyannote/audio/cli/train.py @@ -26,10 +26,11 @@ import hydra from hydra.utils import instantiate +from lightning.pytorch import seed_everything from omegaconf import DictConfig, OmegaConf # from pyannote.audio.core.callback import GraduallyUnfreeze -from pyannote.database import FileFinder, get_protocol +from pyannote.database import FileFinder, registry from pytorch_lightning.callbacks import ( EarlyStopping, LearningRateMonitor, @@ -37,7 +38,6 @@ RichProgressBar, ) from pytorch_lightning.loggers import TensorBoardLogger -from pytorch_lightning.utilities.seed import seed_everything from torch_audiomentations.utils.config import from_dict as get_augmentation from pyannote.audio.core.io import get_torchaudio_info @@ -45,18 +45,21 @@ @hydra.main(config_path="train_config", config_name="config") def train(cfg: DictConfig) -> Optional[float]: - # make sure to set the random seed before the instantiation of Trainer # so that each model initializes with the same weights when using DDP. seed = int(os.environ.get("PL_GLOBAL_SEED", "0")) seed_everything(seed=seed) + # load databases into registry + for database_yml in cfg.registry.split(","): + registry.load_database(database_yml) + # instantiate training protocol with optional preprocessors preprocessors = {"audio": FileFinder(), "torchaudio.info": get_torchaudio_info} if "preprocessor" in cfg: preprocessor = instantiate(cfg.preprocessor) preprocessors[preprocessor.preprocessed_key] = preprocessor - protocol = get_protocol(cfg.protocol, preprocessors=preprocessors) + protocol = registry.get_protocol(cfg.protocol, preprocessors=preprocessors) # instantiate data augmentation augmentation = ( @@ -96,7 +99,11 @@ def configure_optimizers(self): model.configure_optimizers = MethodType(configure_optimizers, model) - callbacks = [RichProgressBar(), LearningRateMonitor(logging_interval="step")] + # avoid creating big log files + callbacks = [ + RichProgressBar(refresh_rate=20, leave=True), + LearningRateMonitor(), + ] if fine_tuning: # TODO: configure layer freezing @@ -108,7 +115,7 @@ def configure_optimizers(self): checkpoint = ModelCheckpoint( monitor=monitor, mode=direction, - save_top_k=None if monitor is None else 5, + save_top_k=None if monitor is None else 1, every_n_epochs=1, save_last=True, save_weights_only=False, @@ -126,6 +133,7 @@ def configure_optimizers(self): patience=100, strict=True, verbose=False, + check_finite=True, ) callbacks.append(early_stopping) diff --git a/pyannote/audio/cli/train_config/config.yaml b/pyannote/audio/cli/train_config/config.yaml index f939f39e1..d5b761cc9 100644 --- a/pyannote/audio/cli/train_config/config.yaml +++ b/pyannote/audio/cli/train_config/config.yaml @@ -1,7 +1,8 @@ +registry: ??? protocol: ??? defaults: - - task: SpeakerSegmentation + - task: SpeakerDiarization - model: PyanNet - optimizer: Adam - scheduler: CosineAnnealingWarmRestarts diff --git a/pyannote/audio/cli/train_config/model/SSeRiouSS.yaml b/pyannote/audio/cli/train_config/model/SSeRiouSS.yaml new file mode 100644 index 000000000..73f7f963a --- /dev/null +++ b/pyannote/audio/cli/train_config/model/SSeRiouSS.yaml @@ -0,0 +1,13 @@ +# @package _group_ +_target_: pyannote.audio.models.segmentation.SSeRiouSS +wav2vec: WAVLM_BASE +wav2vec_layer: -1 +lstm: + hidden_size: 128 + num_layers: 4 + bidirectional: true + monolithic: true + dropout: 0.5 +linear: + hidden_size: 128 + num_layers: 2 diff --git a/pyannote/audio/cli/train_config/optimizer/AdamW.yaml b/pyannote/audio/cli/train_config/optimizer/AdamW.yaml new file mode 100644 index 000000000..f917fb4ed --- /dev/null +++ b/pyannote/audio/cli/train_config/optimizer/AdamW.yaml @@ -0,0 +1,7 @@ +# @package _group_ +_target_: torch.optim.AdamW +lr: 1e-3 +betas: [0.9, 0.999] +eps: 1e-08 +weight_decay: 0.01 +amsgrad: False diff --git a/pyannote/audio/cli/train_config/task/SpeakerSegmentation.yaml b/pyannote/audio/cli/train_config/task/MultiLabelSegmentation.yaml similarity index 56% rename from pyannote/audio/cli/train_config/task/SpeakerSegmentation.yaml rename to pyannote/audio/cli/train_config/task/MultiLabelSegmentation.yaml index b7b6bdc02..aceb8cb22 100644 --- a/pyannote/audio/cli/train_config/task/SpeakerSegmentation.yaml +++ b/pyannote/audio/cli/train_config/task/MultiLabelSegmentation.yaml @@ -1,11 +1,9 @@ # @package _group_ -_target_: pyannote.audio.tasks.Segmentation -duration: 5.0 +_target_: pyannote.audio.tasks.MultiLabelSegmentation +duration: 3.0 warm_up: 0.0 balance: null weight: null batch_size: 32 num_workers: null pin_memory: False -loss: "bce" -vad_loss: "bce" diff --git a/pyannote/audio/cli/train_config/task/SpeakerDiarization.yaml b/pyannote/audio/cli/train_config/task/SpeakerDiarization.yaml new file mode 100644 index 000000000..d9ec04a57 --- /dev/null +++ b/pyannote/audio/cli/train_config/task/SpeakerDiarization.yaml @@ -0,0 +1,8 @@ +# @package _group_ +_target_: pyannote.audio.tasks.SpeakerDiarization +duration: 5.0 +max_speakers_per_chunk: 3 +max_speakers_per_frame: 2 +batch_size: 32 +num_workers: 10 +pin_memory: False diff --git a/pyannote/audio/cli/train_config/trainer/default.yaml b/pyannote/audio/cli/train_config/trainer/default.yaml index eeb5b85b9..ac3a60ff4 100644 --- a/pyannote/audio/cli/train_config/trainer/default.yaml +++ b/pyannote/audio/cli/train_config/trainer/default.yaml @@ -2,46 +2,34 @@ _target_: pytorch_lightning.Trainer accelerator: auto accumulate_grad_batches: 1 -amp_backend: native -auto_lr_find: False -auto_scale_batch_size: False -auto_select_gpus: True -benchmark: False -check_val_every_n_epoch: 1 -detect_anomaly: False +benchmark: null # TODO: automatically set to True when using fixed duration chunks deterministic: False +check_val_every_n_epoch: 1 devices: auto +detect_anomaly: False enable_checkpointing: True enable_model_summary: True enable_progress_bar: True fast_dev_run: False -gpus: null -gradient_clip_val: 0 +gradient_clip_val: null gradient_clip_algorithm: norm -ipus: null limit_predict_batches: 1.0 limit_test_batches: 1.0 limit_train_batches: 1.0 limit_val_batches: 1.0 log_every_n_steps: 50 max_epochs: 1000 -max_steps: null +max_steps: -1 max_time: null min_epochs: 1 min_steps: null -move_metrics_to_cpu: False -multiple_trainloader_mode: max_size_cycle num_nodes: 1 -num_processes: 1 num_sanity_val_steps: 2 overfit_batches: 0.0 precision: 32 profiler: null reload_dataloaders_every_n_epochs: 0 -replace_sampler_ddp: True +use_distributed_sampler: True # TODO: check what this does exactly strategy: null sync_batchnorm: False -tpu_cores: null -track_grad_norm: -1 val_check_interval: 1.0 -weights_save_path: null diff --git a/pyannote/audio/core/inference.py b/pyannote/audio/core/inference.py index b664174b5..dcf21868d 100644 --- a/pyannote/audio/core/inference.py +++ b/pyannote/audio/core/inference.py @@ -1,6 +1,6 @@ # MIT License # -# Copyright (c) 2020-2021 CNRS +# Copyright (c) 2020- CNRS # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal @@ -23,24 +23,30 @@ import math import warnings from pathlib import Path -from typing import Any, Callable, List, Optional, Text, Tuple, Union +from typing import Callable, List, Optional, Text, Tuple, Union import numpy as np import torch +import torch.nn as nn +import torch.nn.functional as F from einops import rearrange from pyannote.core import Segment, SlidingWindow, SlidingWindowFeature from pytorch_lightning.utilities.memory import is_oom_error from pyannote.audio.core.io import AudioFile -from pyannote.audio.core.model import Model +from pyannote.audio.core.model import Model, Specifications from pyannote.audio.core.task import Resolution +from pyannote.audio.utils.multi_task import map_with_specifications from pyannote.audio.utils.permutation import mae_cost_func, permutate -from pyannote.audio.utils.progress import InferenceProgressHook +from pyannote.audio.utils.powerset import Powerset +from pyannote.audio.utils.reproducibility import fix_reproducibility -TaskName = Union[Text, None] +class BaseInference: + pass -class Inference: + +class Inference(BaseInference): """Inference Parameters @@ -50,29 +56,26 @@ class Inference: window : {"sliding", "whole"}, optional Use a "sliding" window and aggregate the corresponding outputs (default) or just one (potentially long) window covering the "whole" file or chunk. - skip_aggregation : bool, optional - Do not aggregate outputs when using "sliding" window. Defaults to False. duration : float, optional Chunk duration, in seconds. Defaults to duration used for training the model. Has no effect when `window` is "whole". step : float, optional Step between consecutive chunks, in seconds. Defaults to warm-up duration when greater than 0s, otherwise 10% of duration. Has no effect when `window` is "whole". - batch_size : int, optional - Batch size. Larger values make inference faster. Defaults to 32. - device : torch.device, optional - Device used for inference. Defaults to `model.device`. - In case `device` and `model.device` are different, model is sent to device. pre_aggregation_hook : callable, optional When a callable is provided, it is applied to the model output, just before aggregation. Takes a (num_chunks, num_frames, dimension) numpy array as input and returns a modified (num_chunks, num_frames, other_dimension) numpy array passed to overlap-add aggregation. - progress_hook : {callable, True, str}, optional - When a callable is provided, it is called everytime a batch is processed - with two integer arguments: - - the number of chunks that have been processed so far - - the total number of chunks - Set to True (or a descriptive string) to display a tqdm progress bar. + skip_aggregation : bool, optional + Do not aggregate outputs when using "sliding" window. Defaults to False. + skip_conversion: bool, optional + In case a task has been trained with `powerset` mode, output is automatically + converted to `multi-label`, unless `skip_conversion` is set to True. + batch_size : int, optional + Batch size. Larger values (should) make inference faster. Defaults to 32. + device : torch.device, optional + Device used for inference. Defaults to `model.device`. + In case `device` and `model.device` are different, model is sent to device. use_auth_token : str, optional When loading a private huggingface.co model, set `use_auth_token` to True or to a string containing your hugginface.co authentication @@ -83,71 +86,92 @@ def __init__( self, model: Union[Model, Text, Path], window: Text = "sliding", - skip_aggregation: bool = False, - device: torch.device = None, duration: float = None, step: float = None, - batch_size: int = 32, pre_aggregation_hook: Callable[[np.ndarray], np.ndarray] = None, - progress_hook: Union[bool, Text, Callable[[int, int], Any]] = False, + skip_aggregation: bool = False, + skip_conversion: bool = False, + device: torch.device = None, + batch_size: int = 32, use_auth_token: Union[Text, None] = None, ): + # ~~~~ model ~~~~~ self.model = ( model if isinstance(model, Model) else Model.from_pretrained( - Path(model), + model, map_location=device, strict=False, use_auth_token=use_auth_token, ) ) - if window not in ["sliding", "whole"]: - raise ValueError('`window` must be "sliding" or "whole".') - - specifications = self.model.specifications - if specifications.resolution == Resolution.FRAME and window == "whole": - warnings.warn( - 'Using "whole" `window` inference with a frame-based model might lead to bad results ' - 'and huge memory consumption: it is recommended to set `window` to "sliding".' - ) - - self.window = window - self.skip_aggregation = skip_aggregation - if device is None: device = self.model.device self.device = device - self.pre_aggregation_hook = pre_aggregation_hook - self.model.eval() self.model.to(self.device) - # chunk duration used during training specifications = self.model.specifications - training_duration = specifications.duration - if duration is None: - duration = training_duration - elif training_duration != duration: + # ~~~~ sliding window ~~~~~ + + if window not in ["sliding", "whole"]: + raise ValueError('`window` must be "sliding" or "whole".') + + if window == "whole" and any( + s.resolution == Resolution.FRAME for s in specifications + ): + warnings.warn( + 'Using "whole" `window` inference with a frame-based model might lead to bad results ' + 'and huge memory consumption: it is recommended to set `window` to "sliding".' + ) + self.window = window + + training_duration = next(iter(specifications)).duration + duration = duration or training_duration + if training_duration != duration: warnings.warn( f"Model was trained with {training_duration:g}s chunks, and you requested " f"{duration:g}s chunks for inference: this might lead to suboptimal results." ) self.duration = duration - self.warm_up = specifications.warm_up + # ~~~~ powerset to multilabel conversion ~~~~ + + self.skip_conversion = skip_conversion + + conversion = list() + for s in specifications: + if s.powerset and not skip_conversion: + c = Powerset(len(s.classes), s.powerset_max_classes) + else: + c = nn.Identity() + conversion.append(c.to(self.device)) + + if isinstance(specifications, Specifications): + self.conversion = conversion[0] + else: + self.conversion = nn.ModuleList(conversion) + + # ~~~~ overlap-add aggregation ~~~~~ + + self.skip_aggregation = skip_aggregation + self.pre_aggregation_hook = pre_aggregation_hook + + self.warm_up = next(iter(specifications)).warm_up # Use that many seconds on the left- and rightmost parts of each chunk # to warm up the model. While the model does process those left- and right-most # parts, only the remaining central part of each chunk is used for aggregating # scores during inference. # step between consecutive chunks - if step is None: - step = 0.1 * self.duration if self.warm_up[0] == 0.0 else self.warm_up[0] + step = step or ( + 0.1 * self.duration if self.warm_up[0] == 0.0 else self.warm_up[0] + ) if step > self.duration: raise ValueError( @@ -159,17 +183,20 @@ def __init__( self.batch_size = batch_size - if callable(progress_hook): - pass - elif isinstance(progress_hook, Text): - progress_hook = InferenceProgressHook(desc=progress_hook) - elif progress_hook: - progress_hook = InferenceProgressHook() - else: - progress_hook = None - self.progress_hook = progress_hook + def to(self, device: torch.device) -> "Inference": + """Send internal model to `device`""" + + if not isinstance(device, torch.device): + raise TypeError( + f"`device` must be an instance of `torch.device`, got `{type(device).__name__}`" + ) - def infer(self, chunks: torch.Tensor) -> np.ndarray: + self.model.to(device) + self.conversion.to(device) + self.device = device + return self + + def infer(self, chunks: torch.Tensor) -> Union[np.ndarray, Tuple[np.ndarray]]: """Forward pass Takes care of sending chunks to right device and outputs back to CPU @@ -181,11 +208,11 @@ def infer(self, chunks: torch.Tensor) -> np.ndarray: Returns ------- - outputs : (batch_size, ...) np.ndarray + outputs : (tuple of) (batch_size, ...) np.ndarray Model output. """ - with torch.no_grad(): + with torch.inference_mode(): try: outputs = self.model(chunks.to(self.device)) except RuntimeError as exception: @@ -197,9 +224,19 @@ def infer(self, chunks: torch.Tensor) -> np.ndarray: else: raise exception - return outputs.cpu().numpy() + def __convert(output: torch.Tensor, conversion: nn.Module, **kwargs): + return conversion(output).cpu().numpy() + + return map_with_specifications( + self.model.specifications, __convert, outputs, self.conversion + ) - def slide(self, waveform: torch.Tensor, sample_rate: int) -> SlidingWindowFeature: + def slide( + self, + waveform: torch.Tensor, + sample_rate: int, + hook: Optional[Callable], + ) -> Union[SlidingWindowFeature, Tuple[SlidingWindowFeature]]: """Slide model on a waveform Parameters @@ -208,26 +245,35 @@ def slide(self, waveform: torch.Tensor, sample_rate: int) -> SlidingWindowFeatur Waveform. sample_rate : int Sample rate. + hook: Optional[Callable] + When a callable is provided, it is called everytime a batch is + processed with two keyword arguments: + - `completed`: the number of chunks that have been processed so far + - `total`: the total number of chunks Returns ------- - output : SlidingWindowFeature + output : (tuple of) SlidingWindowFeature Model output. Shape is (num_chunks, dimension) for chunk-level tasks, and (num_frames, dimension) for frame-level tasks. """ - window_size: int = round(self.duration * sample_rate) + window_size: int = self.model.audio.get_num_samples(self.duration) step_size: int = round(self.step * sample_rate) - num_channels, num_samples = waveform.shape - - specifications = self.model.specifications - resolution = specifications.resolution - introspection = self.model.introspection - if resolution == Resolution.CHUNK: - frames = SlidingWindow(start=0.0, duration=self.duration, step=self.step) - elif resolution == Resolution.FRAME: - frames = introspection.frames - num_frames_per_chunk, dimension = introspection(window_size) + _, num_samples = waveform.shape + + def __frames( + example_output, specifications: Optional[Specifications] = None + ) -> SlidingWindow: + if specifications.resolution == Resolution.CHUNK: + return SlidingWindow(start=0.0, duration=self.duration, step=self.step) + return example_output.frames + + frames: Union[SlidingWindow, Tuple[SlidingWindow]] = map_with_specifications( + self.model.specifications, + __frames, + self.model.example_output, + ) # prepare complete chunks if num_samples >= window_size: @@ -244,100 +290,159 @@ def slide(self, waveform: torch.Tensor, sample_rate: int) -> SlidingWindowFeatur num_samples - window_size ) % step_size > 0 if has_last_chunk: + # pad last chunk with zeros last_chunk: torch.Tensor = waveform[:, num_chunks * step_size :] + _, last_window_size = last_chunk.shape + last_pad = window_size - last_window_size + last_chunk = F.pad(last_chunk, (0, last_pad)) - outputs: Union[List[np.ndarray], np.ndarray] = list() + def __empty_list(**kwargs): + return list() - if self.progress_hook is not None: - self.progress_hook(0, num_chunks + has_last_chunk) + outputs: Union[ + List[np.ndarray], Tuple[List[np.ndarray]] + ] = map_with_specifications(self.model.specifications, __empty_list) + + if hook is not None: + hook(completed=0, total=num_chunks + has_last_chunk) + + def __append_batch(output, batch_output, **kwargs) -> None: + output.append(batch_output) + return # slide over audio chunks in batch for c in np.arange(0, num_chunks, self.batch_size): batch: torch.Tensor = chunks[c : c + self.batch_size] - outputs.append(self.infer(batch)) - if self.progress_hook is not None: - self.progress_hook(c + self.batch_size, num_chunks + has_last_chunk) + + batch_outputs: Union[np.ndarray, Tuple[np.ndarray]] = self.infer(batch) + + _ = map_with_specifications( + self.model.specifications, __append_batch, outputs, batch_outputs + ) + + if hook is not None: + hook(completed=c + self.batch_size, total=num_chunks + has_last_chunk) # process orphan last chunk if has_last_chunk: + last_outputs = self.infer(last_chunk[None]) + + _ = map_with_specifications( + self.model.specifications, __append_batch, outputs, last_outputs + ) + + if hook is not None: + hook( + completed=num_chunks + has_last_chunk, + total=num_chunks + has_last_chunk, + ) - last_output = self.infer(last_chunk[None]) + def __vstack(output: List[np.ndarray], **kwargs) -> np.ndarray: + return np.vstack(output) - if specifications.resolution == Resolution.FRAME: - pad = num_frames_per_chunk - last_output.shape[1] - last_output = np.pad(last_output, ((0, 0), (0, pad), (0, 0))) + outputs: Union[np.ndarray, Tuple[np.ndarray]] = map_with_specifications( + self.model.specifications, __vstack, outputs + ) - outputs.append(last_output) - if self.progress_hook is not None: - self.progress_hook( - num_chunks + has_last_chunk, num_chunks + has_last_chunk + def __aggregate( + outputs: np.ndarray, + frames: SlidingWindow, + specifications: Optional[Specifications] = None, + ) -> SlidingWindowFeature: + # skip aggregation when requested, + # or when model outputs just one vector per chunk + # or when model is permutation-invariant (and not post-processed) + if ( + self.skip_aggregation + or specifications.resolution == Resolution.CHUNK + or ( + specifications.permutation_invariant + and self.pre_aggregation_hook is None + ) + ): + frames = SlidingWindow( + start=0.0, duration=self.duration, step=self.step ) + return SlidingWindowFeature(outputs, frames) + + if self.pre_aggregation_hook is not None: + outputs = self.pre_aggregation_hook(outputs) - outputs = np.vstack(outputs) - - # skip aggregation when requested, - # or when model outputs just one vector per chunk - # or when model is permutation-invariant (and not post-processed) - if ( - self.skip_aggregation - or specifications.resolution == Resolution.CHUNK - or ( - specifications.permutation_invariant - and self.pre_aggregation_hook is None + aggregated = self.aggregate( + SlidingWindowFeature( + outputs, + SlidingWindow(start=0.0, duration=self.duration, step=self.step), + ), + frames=frames, + warm_up=self.warm_up, + hamming=True, + missing=0.0, ) - ): - frames = SlidingWindow(start=0.0, duration=self.duration, step=self.step) - return SlidingWindowFeature(outputs, frames) - - if self.pre_aggregation_hook is not None: - outputs = self.pre_aggregation_hook(outputs) - - aggregated = self.aggregate( - SlidingWindowFeature( - outputs, - SlidingWindow(start=0.0, duration=self.duration, step=self.step), - ), - frames=frames, - warm_up=self.warm_up, - hamming=True, - missing=0.0, - ) - if has_last_chunk: - num_frames = aggregated.data.shape[0] - aggregated.data = aggregated.data[: num_frames - pad, :] + # remove padding that was added to last chunk + if has_last_chunk: + aggregated.data = aggregated.crop( + Segment(0.0, num_samples / sample_rate), mode="loose" + ) - return aggregated + return aggregated - def __call__(self, file: AudioFile) -> Union[SlidingWindowFeature, np.ndarray]: + return map_with_specifications( + self.model.specifications, __aggregate, outputs, frames + ) + + def __call__( + self, file: AudioFile, hook: Optional[Callable] = None + ) -> Union[ + Tuple[Union[SlidingWindowFeature, np.ndarray]], + Union[SlidingWindowFeature, np.ndarray], + ]: """Run inference on a whole file Parameters ---------- file : AudioFile Audio file. + hook : callable, optional + When a callable is provided, it is called everytime a batch is processed + with two keyword arguments: + - `completed`: the number of chunks that have been processed so far + - `total`: the total number of chunks Returns ------- - output : SlidingWindowFeature or np.ndarray + output : (tuple of) SlidingWindowFeature or np.ndarray Model output, as `SlidingWindowFeature` if `window` is set to "sliding" and `np.ndarray` if is set to "whole". """ + fix_reproducibility(self.device) + waveform, sample_rate = self.model.audio(file) if self.window == "sliding": - return self.slide(waveform, sample_rate) + return self.slide(waveform, sample_rate, hook=hook) + + outputs: Union[np.ndarray, Tuple[np.ndarray]] = self.infer(waveform[None]) - return self.infer(waveform[None])[0] + def __first_sample(outputs: np.ndarray, **kwargs) -> np.ndarray: + return outputs[0] + + return map_with_specifications( + self.model.specifications, __first_sample, outputs + ) def crop( self, file: AudioFile, chunk: Union[Segment, List[Segment]], duration: Optional[float] = None, - ) -> Union[SlidingWindowFeature, np.ndarray]: + hook: Optional[Callable] = None, + ) -> Union[ + Tuple[Union[SlidingWindowFeature, np.ndarray]], + Union[SlidingWindowFeature, np.ndarray], + ]: """Run inference on a chunk or a list of chunks Parameters @@ -354,10 +459,15 @@ def crop( Enforce chunk duration (in seconds). This is a hack to avoid rounding errors that may result in a different number of audio samples for two chunks of the same duration. + hook : callable, optional + When a callable is provided, it is called everytime a batch is processed + with two keyword arguments: + - `completed`: the number of chunks that have been processed so far + - `total`: the total number of chunks Returns ------- - output : SlidingWindowFeature or np.ndarray + output : (tuple of) SlidingWindowFeature or np.ndarray Model output, as `SlidingWindowFeature` if `window` is set to "sliding" and `np.ndarray` if is set to "whole". @@ -371,8 +481,9 @@ def crop( >>> inference.crop(file, extended_chunk).crop(chunk_of_interest, returns_data=False) """ - if self.window == "sliding": + fix_reproducibility(self.device) + if self.window == "sliding": if not isinstance(chunk, Segment): start = min(c.start for c in chunk) end = max(c.end for c in chunk) @@ -381,32 +492,37 @@ def crop( waveform, sample_rate = self.model.audio.crop( file, chunk, duration=duration ) - output = self.slide(waveform, sample_rate) - - frames = output.sliding_window - shifted_frames = SlidingWindow( - start=chunk.start, duration=frames.duration, step=frames.step - ) - return SlidingWindowFeature(output.data, shifted_frames) - - elif self.window == "whole": - - if isinstance(chunk, Segment): - waveform, sample_rate = self.model.audio.crop( - file, chunk, duration=duration - ) - else: - waveform = torch.cat( - [self.model.audio.crop(file, c)[0] for c in chunk], dim=1 + outputs: Union[ + SlidingWindowFeature, Tuple[SlidingWindowFeature] + ] = self.slide(waveform, sample_rate, hook=hook) + + def __shift(output: SlidingWindowFeature, **kwargs) -> SlidingWindowFeature: + frames = output.sliding_window + shifted_frames = SlidingWindow( + start=chunk.start, duration=frames.duration, step=frames.step ) + return SlidingWindowFeature(output.data, shifted_frames) - return self.infer(waveform[None])[0] + return map_with_specifications(self.model.specifications, __shift, outputs) + if isinstance(chunk, Segment): + waveform, sample_rate = self.model.audio.crop( + file, chunk, duration=duration + ) else: - raise NotImplementedError( - f"Unsupported window type '{self.window}': should be 'sliding' or 'whole'." + waveform = torch.cat( + [self.model.audio.crop(file, c)[0] for c in chunk], dim=1 ) + outputs: Union[np.ndarray, Tuple[np.ndarray]] = self.infer(waveform[None]) + + def __first_sample(outputs: np.ndarray, **kwargs) -> np.ndarray: + return outputs[0] + + return map_with_specifications( + self.model.specifications, __first_sample, outputs + ) + @staticmethod def aggregate( scores: SlidingWindowFeature, @@ -648,7 +764,6 @@ def always_match(this: np.ndarray, that: np.ndarray, cost: float): stitches = [] for C, (chunk, activation) in enumerate(activations): - local_stitch = np.NAN * np.zeros( (sum(lookahead) + 1, num_frames, num_classes) ) @@ -656,7 +771,6 @@ def always_match(this: np.ndarray, that: np.ndarray, cost: float): for c in range( max(0, C - lookahead[0]), min(num_chunks, C + lookahead[1] + 1) ): - # extract common temporal support shift = round((C - c) * num_frames * chunks.step / chunks.duration) @@ -677,7 +791,6 @@ def always_match(this: np.ndarray, that: np.ndarray, cost: float): ) for this, that in enumerate(permutation): - # only stitch under certain condiditions matching = (c == C) or ( match_func( diff --git a/pyannote/audio/core/io.py b/pyannote/audio/core/io.py index 49175f531..0a44e75ea 100644 --- a/pyannote/audio/core/io.py +++ b/pyannote/audio/core/io.py @@ -28,6 +28,7 @@ """ import math +import random import warnings from io import IOBase from pathlib import Path @@ -79,12 +80,14 @@ class Audio: ---------- sample_rate: int, optional Target sampling rate. Defaults to using native sampling rate. - mono : int, optional - Convert multi-channel to mono. Defaults to True. + mono : {'random', 'downmix'}, optional + In case of multi-channel audio, convert to single-channel audio + using one of the following strategies: select one channel at + 'random' or 'downmix' by averaging all channels. Usage ----- - >>> audio = Audio(sample_rate=16000, mono=True) + >>> audio = Audio(sample_rate=16000, mono='downmix') >>> waveform, sample_rate = audio({"audio": "/path/to/audio.wav"}) >>> assert sample_rate == 16000 >>> sample_rate = 44100 @@ -147,7 +150,6 @@ def validate_file(file: AudioFile) -> Mapping: raise ValueError(AudioFileDocString) if "waveform" in file: - waveform: Union[np.ndarray, Tensor] = file["waveform"] if len(waveform.shape) != 2 or waveform.shape[0] > waveform.shape[1]: raise ValueError( @@ -163,7 +165,6 @@ def validate_file(file: AudioFile) -> Mapping: file.setdefault("uri", "waveform") elif "audio" in file: - if isinstance(file["audio"], IOBase): return file @@ -174,15 +175,13 @@ def validate_file(file: AudioFile) -> Mapping: file.setdefault("uri", path.stem) else: - raise ValueError( "Neither 'waveform' nor 'audio' is available for this file." ) return file - def __init__(self, sample_rate=None, mono=True): - + def __init__(self, sample_rate=None, mono=None): super().__init__() self.sample_rate = sample_rate self.mono = mono @@ -206,8 +205,13 @@ def downmix_and_resample(self, waveform: Tensor, sample_rate: int) -> Tensor: """ # downmix to mono - if self.mono and waveform.shape[0] > 1: - waveform = waveform.mean(dim=0, keepdim=True) + num_channels = waveform.shape[0] + if num_channels > 1: + if self.mono == "random": + channel = random.randint(0, num_channels - 1) + waveform = waveform[channel : channel + 1] + elif self.mono == "downmix": + waveform = waveform.mean(dim=0, keepdim=True) # resample if (self.sample_rate is not None) and (self.sample_rate != sample_rate): @@ -249,6 +253,18 @@ def get_duration(self, file: AudioFile) -> float: return frames / sample_rate + def get_num_samples(self, duration: float, sample_rate: int = None) -> int: + """Deterministic number of samples from duration and sample rate""" + + sample_rate = sample_rate or self.sample_rate + + if sample_rate is None: + raise ValueError( + "`sample_rate` must be provided to compute number of samples." + ) + + return math.floor(duration * sample_rate) + def __call__(self, file: AudioFile) -> Tuple[Tensor, int]: """Obtain waveform @@ -277,10 +293,14 @@ def __call__(self, file: AudioFile) -> Tuple[Tensor, int]: elif "audio" in file: waveform, sample_rate = torchaudio.load(file["audio"]) + # rewind if needed + if isinstance(file["audio"], IOBase): + file["audio"].seek(0) + channel = file.get("channel", None) if channel is not None: - waveform = waveform[channel - 1 : channel] + waveform = waveform[channel : channel + 1] return self.downmix_and_resample(waveform, sample_rate) @@ -347,7 +367,6 @@ def crop( num_frames = end_frame - start_frame if mode == "raise": - if num_frames > frames: raise ValueError( f"requested fixed duration ({duration:6f}s, or {num_frames:d} frames) is longer " @@ -384,10 +403,10 @@ def crop( data, _ = torchaudio.load( file["audio"], frame_offset=start_frame, num_frames=num_frames ) + # rewind if needed if isinstance(file["audio"], IOBase): file["audio"].seek(0) except RuntimeError: - if isinstance(file["audio"], IOBase): msg = "torchaudio failed to seek-and-read in file-like object." raise RuntimeError(msg) @@ -408,7 +427,7 @@ def crop( file["sample_rate"] = sample_rate if channel is not None: - data = data[channel - 1 : channel, :] + data = data[channel : channel + 1, :] # pad with zeros if mode == "pad": diff --git a/pyannote/audio/core/model.py b/pyannote/audio/core/model.py index ddf4b2795..bedb7f6c4 100644 --- a/pyannote/audio/core/model.py +++ b/pyannote/audio/core/model.py @@ -1,6 +1,6 @@ # MIT License # -# Copyright (c) 2020-2021 CNRS +# Copyright (c) 2020- CNRS # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal @@ -24,6 +24,8 @@ import os import warnings +from dataclasses import dataclass +from functools import cached_property from importlib import import_module from pathlib import Path from typing import Any, Dict, List, Optional, Text, Tuple, Union @@ -33,16 +35,24 @@ import torch import torch.nn as nn import torch.optim -from huggingface_hub import cached_download, hf_hub_url +from huggingface_hub import hf_hub_download +from huggingface_hub.utils import RepositoryNotFoundError +from lightning_fabric.utilities.cloud_io import _load as pl_load from pyannote.core import SlidingWindow -from pytorch_lightning.utilities.cloud_io import load as pl_load from pytorch_lightning.utilities.model_summary import ModelSummary -from semver import VersionInfo from torch.utils.data import DataLoader from pyannote.audio import __version__ from pyannote.audio.core.io import Audio -from pyannote.audio.core.task import Problem, Resolution, Specifications, Task +from pyannote.audio.core.task import ( + Problem, + Resolution, + Specifications, + Task, + UnknownSpecificationsError, +) +from pyannote.audio.utils.multi_task import map_with_specifications +from pyannote.audio.utils.version import check_version CACHE_DIR = os.getenv( "PYANNOTE_CACHE", @@ -52,195 +62,16 @@ HF_LIGHTNING_CONFIG_NAME = "config.yaml" +# NOTE: needed to backward compatibility to load models trained before pyannote.audio 3.x class Introspection: - """Model introspection - - Parameters - ---------- - min_num_samples: int - Minimum number of input samples - min_num_frames: int - Corresponding minimum number of output frames - inc_num_samples: int - Number of input samples leading to an increase of number of output frames - inc_num_frames: int - Corresponding increase in number of output frames - dimension: int - Output dimension - sample_rate: int - Expected input sample rate - - Usage - ----- - >>> introspection = Introspection.from_model(model) - >>> isinstance(introspection.frames, SlidingWindow) - >>> num_samples = 16000 # 1s at 16kHz - >>> num_frames, dimension = introspection(num_samples) - """ - - def __init__( - self, - min_num_samples: int, - min_num_frames: int, - inc_num_samples: int, - inc_num_frames: int, - dimension: int, - sample_rate: int, - ): - super().__init__() - self.min_num_samples = min_num_samples - self.min_num_frames = min_num_frames - self.inc_num_samples = inc_num_samples - self.inc_num_frames = inc_num_frames - self.dimension = dimension - self.sample_rate = sample_rate + pass - @classmethod - def from_model(cls, model: "Model", task: str = None) -> Introspection: - - specifications = model.specifications - if task is not None: - specifications = specifications[task] - - example_input_array = model.example_input_array - batch_size, num_channels, num_samples = example_input_array.shape - example_input_array = torch.randn( - (batch_size, num_channels, num_samples), - dtype=example_input_array.dtype, - layout=example_input_array.layout, - device=example_input_array.device, - requires_grad=False, - ) - # dichotomic search of "min_num_samples" - lower, upper, min_num_samples = 1, num_samples, None - while True: - num_samples = (lower + upper) // 2 - try: - with torch.no_grad(): - frames = model(example_input_array[:, :, :num_samples]) - if task is not None: - frames = frames[task] - except Exception: - lower = num_samples - else: - min_num_samples = num_samples - if specifications.resolution == Resolution.FRAME: - _, min_num_frames, dimension = frames.shape - elif specifications.resolution == Resolution.CHUNK: - _, dimension = frames.shape - else: - # should never happen - pass - upper = num_samples - - if lower + 1 == upper: - break - - # if "min_num_samples" is still None at this point, it means that - # the forward pass always failed and raised an exception. most likely, - # it means that there is a problem with the model definition. - # we try again without catching the exception to help the end user debug - # their model - if min_num_samples is None: - frames = model(example_input_array) - - # corner case for chunk-level tasks - if specifications.resolution == Resolution.CHUNK: - return cls( - min_num_samples=min_num_samples, - min_num_frames=1, - inc_num_samples=0, - inc_num_frames=0, - dimension=dimension, - sample_rate=model.hparams.sample_rate, - ) - - # search reasonable upper bound for "inc_num_samples" - while True: - num_samples = 2 * min_num_samples - example_input_array = torch.randn( - (batch_size, num_channels, num_samples), - dtype=example_input_array.dtype, - layout=example_input_array.layout, - device=example_input_array.device, - requires_grad=False, - ) - with torch.no_grad(): - frames = model(example_input_array) - if task is not None: - frames = frames[task] - num_frames = frames.shape[1] - if num_frames > min_num_frames: - break - - # dichotomic search of "inc_num_samples" - lower, upper = min_num_samples, num_samples - while True: - num_samples = (lower + upper) // 2 - example_input_array = torch.randn( - (batch_size, num_channels, num_samples), - dtype=example_input_array.dtype, - layout=example_input_array.layout, - device=example_input_array.device, - requires_grad=False, - ) - with torch.no_grad(): - frames = model(example_input_array) - if task is not None: - frames = frames[task] - num_frames = frames.shape[1] - if num_frames > min_num_frames: - inc_num_frames = num_frames - min_num_frames - inc_num_samples = num_samples - min_num_samples - upper = num_samples - else: - lower = num_samples - - if lower + 1 == upper: - break - - return cls( - min_num_samples=min_num_samples, - min_num_frames=min_num_frames, - inc_num_samples=inc_num_samples, - inc_num_frames=inc_num_frames, - dimension=dimension, - sample_rate=model.hparams.sample_rate, - ) - - def __call__(self, num_samples: int) -> Tuple[int, int]: - """Predict output shape, given number of input samples - - Parameters - ---------- - num_samples : int - Number of input samples. - - Returns - ------- - num_frames : int - Number of output frames - dimension : int - Dimension of output frames - """ - - if num_samples < self.min_num_samples: - return 0, self.dimension - - return ( - self.min_num_frames - + self.inc_num_frames - * ((num_samples - self.min_num_samples + 1) // self.inc_num_samples), - self.dimension, - ) - - @property - def frames(self) -> SlidingWindow: - # HACK to support model trained before 'sample_rate' was an Introspection attribute - sample_rate = getattr(self, "sample_rate", 16000) - step = (self.inc_num_samples / self.inc_num_frames) / sample_rate - return SlidingWindow(start=0.0, step=step, duration=step) +@dataclass +class Output: + num_frames: int + dimension: int + frames: SlidingWindow class Model(pl.LightningModule): @@ -271,42 +102,69 @@ def __init__( self.save_hyperparameters("sample_rate", "num_channels") self.task = task - self.audio = Audio( - sample_rate=self.hparams.sample_rate, mono=self.hparams.num_channels == 1 - ) - - @property - def example_input_array(self) -> torch.Tensor: - batch_size = 3 if self.task is None else self.task.batch_size - duration = 2.0 if self.task is None else self.task.duration - - return torch.randn( - ( - batch_size, - self.hparams.num_channels, - int(self.hparams.sample_rate * duration), - ), - device=self.device, - ) + self.audio = Audio(sample_rate=self.hparams.sample_rate, mono="downmix") @property - def task(self): + def task(self) -> Task: return self._task @task.setter - def task(self, task): - self._task = task - del self.introspection + def task(self, task: Task): + # reset (cached) properties when task changes del self.specifications + try: + del self.example_output + except AttributeError: + pass + self._task = task + + def build(self): + # use this method to add task-dependent layers to the model + # (e.g. the final classification and activation layers) + pass @property - def specifications(self): - if self.task is not None: - return self.task.specifications - return self._specifications + def specifications(self) -> Union[Specifications, Tuple[Specifications]]: + if self.task is None: + try: + specifications = self._specifications + + except AttributeError as e: + raise UnknownSpecificationsError( + "Model specifications are not available because it has not been assigned a task yet. " + "Use `model.task = ...` to assign a task to the model." + ) from e + + else: + try: + specifications = self.task.specifications + + except AttributeError as e: + raise UnknownSpecificationsError( + "Task specifications are not available. This is most likely because they depend on " + "the content of the training subset. Use `model.task.setup()` to go over the training " + "subset and fix this, or let lightning trainer do that for you in `trainer.fit(model)`." + ) from e + + return specifications @specifications.setter - def specifications(self, specifications): + def specifications( + self, specifications: Union[Specifications, Tuple[Specifications]] + ): + if not isinstance(specifications, (Specifications, tuple)): + raise ValueError( + "Only regular specifications or tuple of specifications are supported." + ) + + durations = set(s.duration for s in specifications) + if len(durations) > 1: + raise ValueError("All tasks must share the same (maximum) duration.") + + min_durations = set(s.min_duration for s in specifications) + if len(min_durations) > 1: + raise ValueError("All tasks must share the same minimum duration.") + self._specifications = specifications @specifications.deleter @@ -314,39 +172,54 @@ def specifications(self): if hasattr(self, "_specifications"): del self._specifications - def build(self): - # use this method to add task-dependent layers to the model - # (e.g. the final classification and activation layers) - pass + def __example_input_array(self, duration: Optional[float] = None) -> torch.Tensor: + duration = duration or next(iter(self.specifications)).duration + return torch.randn( + ( + 1, + self.hparams.num_channels, + self.audio.get_num_samples(duration), + ), + device=self.device, + ) @property - def introspection(self) -> Introspection: - """Introspection - - Returns - ------- - introspection: Introspection - Model introspection - """ - - if not hasattr(self, "_introspection"): - self._introspection = Introspection.from_model(self) - - return self._introspection + def example_input_array(self) -> torch.Tensor: + return self.__example_input_array() + + @cached_property + def example_output(self) -> Union[Output, Tuple[Output]]: + """Example output""" + example_input_array = self.__example_input_array() + with torch.inference_mode(): + example_output = self(example_input_array) + + def __example_output( + example_output: torch.Tensor, + specifications: Specifications = None, + ) -> Output: + if specifications.resolution == Resolution.FRAME: + _, num_frames, dimension = example_output.shape + frame_duration = specifications.duration / num_frames + frames = SlidingWindow(step=frame_duration, duration=frame_duration) + else: + _, dimension = example_output.shape + num_frames = None + frames = None - @introspection.setter - def introspection(self, introspection): - self._introspection = introspection + return Output( + num_frames=num_frames, + dimension=dimension, + frames=frames, + ) - @introspection.deleter - def introspection(self): - if hasattr(self, "_introspection"): - del self._introspection + return map_with_specifications( + self.specifications, __example_output, example_output + ) def setup(self, stage=None): - if stage == "fit": - self.task.setup() + self.task.setup_metadata() # list of layers before adding task-dependent layers before = set((name, id(module)) for name, module in self.named_modules()) @@ -387,8 +260,8 @@ def setup(self, stage=None): # setup custom validation metrics self.task.setup_validation_metric() - # this is to make sure introspection is performed here, once and for all - _ = self.introspection + # cache for later (and to avoid later CUDA error with multiprocessing) + _ = self.example_output # list of layers after adding task-dependent layers after = set((name, id(module)) for name, module in self.named_modules()) @@ -397,7 +270,6 @@ def setup(self, stage=None): self.task_dependent = list(name for name, _ in after - before) def on_save_checkpoint(self, checkpoint): - # put everything pyannote.audio-specific under pyannote.audio # to avoid any future conflicts with pytorch-lightning updates checkpoint["pyannote.audio"] = { @@ -409,82 +281,44 @@ def on_save_checkpoint(self, checkpoint): "module": self.__class__.__module__, "class": self.__class__.__name__, }, - "introspection": self.introspection, "specifications": self.specifications, } - @staticmethod - def check_version(library: Text, theirs: Text, mine: Text): - theirs = VersionInfo.parse(theirs) - mine = VersionInfo.parse(mine) - if theirs.major != mine.major: - warnings.warn( - f"Model was trained with {library} {theirs}, yours is {mine}. " - f"Bad things will probably happen unless you update {library} to {theirs.major}.x." - ) - if theirs.minor > mine.minor: - warnings.warn( - f"Model was trained with {library} {theirs}, yours is {mine}. " - f"This should be OK but you might want to update {library}." - ) - def on_load_checkpoint(self, checkpoint: Dict[str, Any]): - - self.check_version( + check_version( "pyannote.audio", checkpoint["pyannote.audio"]["versions"]["pyannote.audio"], __version__, + what="Model", ) - self.check_version( + check_version( "torch", checkpoint["pyannote.audio"]["versions"]["torch"], torch.__version__, + what="Model", ) - self.check_version( - "pytorch-lightning", checkpoint["pytorch-lightning_version"], pl.__version__ + check_version( + "pytorch-lightning", + checkpoint["pytorch-lightning_version"], + pl.__version__, + what="Model", ) self.specifications = checkpoint["pyannote.audio"]["specifications"] + # add task-dependent (e.g. final classifier) layers self.setup() - self.introspection = checkpoint["pyannote.audio"]["introspection"] - - def forward(self, waveforms: torch.Tensor) -> torch.Tensor: + def forward( + self, waveforms: torch.Tensor, **kwargs + ) -> Union[torch.Tensor, Tuple[torch.Tensor]]: msg = "Class {self.__class__.__name__} should define a `forward` method." raise NotImplementedError(msg) - def helper_default_activation(self, specifications: Specifications) -> nn.Module: - """Helper function for default_activation - - Parameters - ---------- - specifications: Specifications - Task specification. - - Returns - ------- - activation : nn.Module - Default activation function. - """ - - if specifications.problem == Problem.BINARY_CLASSIFICATION: - return nn.Sigmoid() - - elif specifications.problem == Problem.MONO_LABEL_CLASSIFICATION: - return nn.LogSoftmax(dim=-1) - - elif specifications.problem == Problem.MULTI_LABEL_CLASSIFICATION: - return nn.Sigmoid() - - else: - msg = "TODO: implement default activation for other types of problems" - raise NotImplementedError(msg) - # convenience function to automate the choice of the final activation function - def default_activation(self) -> nn.Module: + def default_activation(self) -> Union[nn.Module, Tuple[nn.Module]]: """Guess default activation function according to task specification * sigmoid for binary classification @@ -493,10 +327,25 @@ def default_activation(self) -> nn.Module: Returns ------- - activation : nn.Module + activation : (tuple of) nn.Module Activation. """ - return self.helper_default_activation(self.specifications) + + def __default_activation(specifications: Specifications = None) -> nn.Module: + if specifications.problem == Problem.BINARY_CLASSIFICATION: + return nn.Sigmoid() + + elif specifications.problem == Problem.MONO_LABEL_CLASSIFICATION: + return nn.LogSoftmax(dim=-1) + + elif specifications.problem == Problem.MULTI_LABEL_CLASSIFICATION: + return nn.Sigmoid() + + else: + msg = "TODO: implement default activation for other types of problems" + raise NotImplementedError(msg) + + return map_with_specifications(self.specifications, __default_activation) # training data logic is delegated to the task because the # model does not really need to know how it is being used. @@ -518,15 +367,10 @@ def val_dataloader(self) -> DataLoader: def validation_step(self, batch, batch_idx): return self.task.validation_step(batch, batch_idx) - def validation_epoch_end(self, outputs): - return self.task.validation_epoch_end(outputs) - def configure_optimizers(self): return torch.optim.Adam(self.parameters(), lr=1e-3) - def _helper_up_to( - self, module_name: Text, requires_grad: bool = False - ) -> List[Text]: + def __up_to(self, module_name: Text, requires_grad: bool = False) -> List[Text]: """Helper function for freeze_up_to and unfreeze_up_to""" tokens = module_name.split(".") @@ -583,7 +427,7 @@ def freeze_up_to(self, module_name: Text) -> List[Text]: If your model does not follow a sequential structure, you might want to use freeze_by_name for more control. """ - return self._helper_up_to(module_name, requires_grad=False) + return self.__up_to(module_name, requires_grad=False) def unfreeze_up_to(self, module_name: Text) -> List[Text]: """Unfreeze model up to specific module @@ -608,9 +452,9 @@ def unfreeze_up_to(self, module_name: Text) -> List[Text]: If your model does not follow a sequential structure, you might want to use freeze_by_name for more control. """ - return self._helper_up_to(module_name, requires_grad=True) + return self.__up_to(module_name, requires_grad=True) - def _helper_by_name( + def __by_name( self, modules: Union[List[Text], Text], recurse: bool = True, @@ -625,7 +469,6 @@ def _helper_by_name( modules = [modules] for name, module in ModelSummary(self, max_depth=-1).named_modules: - if name not in modules: continue @@ -667,7 +510,7 @@ def freeze_by_name( ValueError if at least one of `modules` does not exist. """ - return self._helper_by_name( + return self.__by_name( modules, recurse=recurse, requires_grad=False, @@ -698,7 +541,7 @@ def unfreeze_by_name( ValueError if at least one of `modules` does not exist. """ - return self._helper_by_name(modules, recurse=recurse, requires_grad=True) + return self.__by_name(modules, recurse=recurse, requires_grad=True) @classmethod def from_pretrained( @@ -777,32 +620,61 @@ def from_pretrained( model_id = checkpoint revision = None - url = hf_hub_url( - model_id, filename=HF_PYTORCH_WEIGHTS_NAME, revision=revision - ) - path_for_pl = cached_download( - url=url, - library_name="pyannote", - library_version=__version__, - cache_dir=cache_dir, - use_auth_token=use_auth_token, - ) + try: + path_for_pl = hf_hub_download( + model_id, + HF_PYTORCH_WEIGHTS_NAME, + repo_type="model", + revision=revision, + library_name="pyannote", + library_version=__version__, + cache_dir=cache_dir, + # force_download=False, + # proxies=None, + # etag_timeout=10, + # resume_download=False, + use_auth_token=use_auth_token, + # local_files_only=False, + # legacy_cache_layout=False, + ) + except RepositoryNotFoundError: + print( + f""" +Could not download '{model_id}' model. +It might be because the model is private or gated so make +sure to authenticate. Visit https://hf.co/settings/tokens to +create your access token and retry with: + + >>> Model.from_pretrained('{model_id}', + ... use_auth_token=YOUR_AUTH_TOKEN) + +If this still does not work, it might be because the model is gated: +visit https://hf.co/{model_id} to accept the user conditions.""" + ) + return None # HACK Huggingface download counters rely on config.yaml # HACK Therefore we download config.yaml even though we # HACK do not use it. Fails silently in case model does not # HACK have a config.yaml file. try: - config_url = hf_hub_url( - model_id, filename=HF_LIGHTNING_CONFIG_NAME, revision=revision - ) - _ = cached_download( - url=config_url, + _ = hf_hub_download( + model_id, + HF_LIGHTNING_CONFIG_NAME, + repo_type="model", + revision=revision, library_name="pyannote", library_version=__version__, cache_dir=cache_dir, + # force_download=False, + # proxies=None, + # etag_timeout=10, + # resume_download=False, use_auth_token=use_auth_token, + # local_files_only=False, + # legacy_cache_layout=False, ) + except Exception: pass diff --git a/pyannote/audio/core/pipeline.py b/pyannote/audio/core/pipeline.py index 90daf297c..f844d584f 100644 --- a/pyannote/audio/core/pipeline.py +++ b/pyannote/audio/core/pipeline.py @@ -22,21 +22,27 @@ import os import warnings +from collections import OrderedDict from collections.abc import Iterator from functools import partial from pathlib import Path -from typing import Callable, List, Optional, Text, Union +from typing import Callable, Dict, List, Optional, Text, Union +import torch import yaml -from huggingface_hub import cached_download, hf_hub_url - -from pyannote.audio import Audio, __version__ -from pyannote.audio.core.io import AudioFile -from pyannote.audio.core.model import CACHE_DIR +from huggingface_hub import hf_hub_download +from huggingface_hub.utils import RepositoryNotFoundError from pyannote.core.utils.helper import get_class_by_name from pyannote.database import FileFinder, ProtocolFile from pyannote.pipeline import Pipeline as _Pipeline +from pyannote.audio import Audio, __version__ +from pyannote.audio.core.inference import BaseInference +from pyannote.audio.core.io import AudioFile +from pyannote.audio.core.model import CACHE_DIR, Model +from pyannote.audio.utils.reproducibility import fix_reproducibility +from pyannote.audio.utils.version import check_version + PIPELINE_PARAMS_NAME = "config.yaml" @@ -77,25 +83,57 @@ def from_pretrained( else: model_id = checkpoint_path revision = None - url = hf_hub_url(model_id, filename=PIPELINE_PARAMS_NAME, revision=revision) - - config_yml = cached_download( - url=url, - library_name="pyannote", - library_version=__version__, - cache_dir=cache_dir, - use_auth_token=use_auth_token, - ) + + try: + config_yml = hf_hub_download( + model_id, + PIPELINE_PARAMS_NAME, + repo_type="model", + revision=revision, + library_name="pyannote", + library_version=__version__, + cache_dir=cache_dir, + # force_download=False, + # proxies=None, + # etag_timeout=10, + # resume_download=False, + use_auth_token=use_auth_token, + # local_files_only=False, + # legacy_cache_layout=False, + ) + + except RepositoryNotFoundError: + print( + f""" +Could not download '{model_id}' pipeline. +It might be because the pipeline is private or gated so make +sure to authenticate. Visit https://hf.co/settings/tokens to +create your access token and retry with: + + >>> Pipeline.from_pretrained('{model_id}', + ... use_auth_token=YOUR_AUTH_TOKEN) + +If this still does not work, it might be because the pipeline is gated: +visit https://hf.co/{model_id} to accept the user conditions.""" + ) + return None with open(config_yml, "r") as fp: config = yaml.load(fp, Loader=yaml.SafeLoader) + if "version" in config: + check_version( + "pyannote.audio", config["version"], __version__, what="Pipeline" + ) + # initialize pipeline pipeline_name = config["pipeline"]["name"] Klass = get_class_by_name( pipeline_name, default_module_name="pyannote.pipeline.blocks" ) - pipeline = Klass(**config["pipeline"].get("params", {})) + params = config["pipeline"].get("params", {}) + params.setdefault("use_auth_token", use_auth_token) + pipeline = Klass(**params) # freeze parameters if "freeze" in config: @@ -111,7 +149,6 @@ def from_pretrained( if "preprocessors" in config: preprocessors = {} for key, preprocessor in config.get("preprocessors", {}).items(): - # preprocessors: # key: # name: package.module.ClassName @@ -139,22 +176,98 @@ def from_pretrained( pipeline.preprocessors = preprocessors + # send pipeline to specified device + if "device" in config: + device = torch.device(config["device"]) + try: + pipeline.to(device) + except RuntimeError as e: + print(e) + return pipeline - @staticmethod - def setup_hook(file: AudioFile, hook: Optional[Callable] = None) -> Callable: + def __init__(self): + super().__init__() + self._models: Dict[str, Model] = OrderedDict() + self._inferences: Dict[str, BaseInference] = OrderedDict() + + def __getattr__(self, name): + """(Advanced) attribute getter + + Adds support for Model and Inference attributes, + which are iterated over by Pipeline.to() method. + + See pyannote.pipeline.Pipeline.__getattr__. + """ + + if "_models" in self.__dict__: + _models = self.__dict__["_models"] + if name in _models: + return _models[name] + + if "_inferences" in self.__dict__: + _inferences = self.__dict__["_inferences"] + if name in _inferences: + return _inferences[name] + + return super().__getattr__(name) + + def __setattr__(self, name, value): + """(Advanced) attribute setter + + Adds support for Model and Inference attributes, + which are iterated over by Pipeline.to() method. + + See pyannote.pipeline.Pipeline.__setattr__. + """ + + def remove_from(*dicts): + for d in dicts: + if name in d: + del d[name] + + _parameters = self.__dict__.get("_parameters") + _instantiated = self.__dict__.get("_instantiated") + _pipelines = self.__dict__.get("_pipelines") + _models = self.__dict__.get("_models") + _inferences = self.__dict__.get("_inferences") + + if isinstance(value, Model): + if _models is None: + msg = "cannot assign models before Pipeline.__init__() call" + raise AttributeError(msg) + remove_from( + self.__dict__, _inferences, _parameters, _instantiated, _pipelines + ) + _models[name] = value + return + + if isinstance(value, BaseInference): + if _inferences is None: + msg = "cannot assign inferences before Pipeline.__init__() call" + raise AttributeError(msg) + remove_from(self.__dict__, _models, _parameters, _instantiated, _pipelines) + _inferences[name] = value + return + + super().__setattr__(name, value) - if hook is None: + def __delattr__(self, name): + if name in self._models: + del self._models[name] - def hook(*args, **kwargs): - return + elif name in self._inferences: + del self._inferences[name] - hook.missing = True else: - hook = partial(hook, file=file) - hook.missing = False + super().__delattr__(name) - return hook + @staticmethod + def setup_hook(file: AudioFile, hook: Optional[Callable] = None) -> Callable: + def noop(*args, **kwargs): + return + + return partial(hook or noop, file=file) def default_parameters(self): raise NotImplementedError() @@ -181,6 +294,8 @@ def classes(self) -> Union[List, Iterator]: raise NotImplementedError() def __call__(self, file: AudioFile, **kwargs): + fix_reproducibility(getattr(self, "device", torch.device("cpu"))) + if not self.instantiated: # instantiate with default parameters when available try: @@ -208,3 +323,25 @@ def __call__(self, file: AudioFile, **kwargs): file = ProtocolFile(file, lazy=self.preprocessors) return self.apply(file, **kwargs) + + def to(self, device: torch.device): + """Send pipeline to `device`""" + + if not isinstance(device, torch.device): + raise TypeError( + f"`device` must be an instance of `torch.device`, got `{type(device).__name__}`" + ) + + for _, pipeline in self._pipelines.items(): + if hasattr(pipeline, "to"): + _ = pipeline.to(device) + + for _, model in self._models.items(): + _ = model.to(device) + + for _, inference in self._inferences.items(): + _ = inference.to(device) + + self.device = device + + return self diff --git a/pyannote/audio/core/task.py b/pyannote/audio/core/task.py index fb42fb7cc..1edfbc35c 100644 --- a/pyannote/audio/core/task.py +++ b/pyannote/audio/core/task.py @@ -23,29 +23,23 @@ from __future__ import annotations -from functools import partial - -try: - from functools import cached_property -except ImportError: - from backports.cached_property import cached_property - import multiprocessing import sys import warnings from dataclasses import dataclass from enum import Enum +from functools import cached_property, partial from numbers import Number -from typing import Dict, List, Optional, Sequence, Text, Tuple, Union +from typing import Dict, List, Literal, Optional, Sequence, Text, Tuple, Union import pytorch_lightning as pl +import scipy.special import torch from pyannote.database import Protocol from torch.utils.data import DataLoader, Dataset, IterableDataset from torch_audiomentations import Identity from torch_audiomentations.core.transforms_interface import BaseWaveformTransform from torchmetrics import Metric, MetricCollection -from typing_extensions import Literal from pyannote.audio.utils.loss import binary_cross_entropy, nll_loss from pyannote.audio.utils.protocol import check_protocol @@ -69,14 +63,20 @@ class Resolution(Enum): CHUNK = 2 # model outputs just one vector for the whole chunk +class UnknownSpecificationsError(Exception): + pass + + @dataclass class Specifications: problem: Problem resolution: Resolution - # chunk duration in seconds. - # use None for variable-length chunks - duration: Optional[float] = None + # (maximum) chunk duration in seconds + duration: float + + # (for variable-duration tasks only) minimum chunk duration in seconds + min_duration: Optional[float] = None # use that many seconds on the left- and rightmost parts of each chunk # to warm up the model. This is mostly useful for segmentation tasks. @@ -89,9 +89,43 @@ class Specifications: # (for classification tasks only) list of classes classes: Optional[List[Text]] = None + # (for powerset only) max number of simultaneous classes + # (n choose k with k <= powerset_max_classes) + powerset_max_classes: Optional[int] = None + # whether classes are permutation-invariant (e.g. diarization) permutation_invariant: bool = False + @cached_property + def powerset(self) -> bool: + if self.powerset_max_classes is None: + return False + + if self.problem != Problem.MONO_LABEL_CLASSIFICATION: + raise ValueError( + "`powerset_max_classes` only makes sense with multi-class classification problems." + ) + + return True + + @cached_property + def num_powerset_classes(self) -> int: + # compute number of subsets of size at most "powerset_max_classes" + # e.g. with len(classes) = 3 and powerset_max_classes = 2: + # {}, {0}, {1}, {2}, {0, 1}, {0, 2}, {1, 2} + return int( + sum( + scipy.special.binom(len(self.classes), i) + for i in range(0, self.powerset_max_classes + 1) + ) + ) + + def __len__(self): + return 1 + + def __iter__(self): + yield self + class TrainDataset(IterableDataset): def __init__(self, task: Task): @@ -165,7 +199,7 @@ class Task(pl.LightningDataModule): Attributes ---------- - specifications : Specifications or dict of Specifications + specifications : Specifications or tuple of Specifications Task specifications (available after `Task.setup` has been called.) """ @@ -184,7 +218,10 @@ def __init__( super().__init__() # dataset - self.protocol, self.has_validation = check_protocol(protocol) + self.protocol, checks = check_protocol(protocol) + self.has_validation = checks["has_validation"] + self.has_scope = checks["has_scope"] + self.has_classes = checks["has_classes"] # batching self.duration = duration @@ -231,7 +268,28 @@ def prepare_data(self): """ pass - def setup(self, stage: Optional[str] = None): + @property + def specifications(self) -> Union[Specifications, Tuple[Specifications]]: + # setup metadata on-demand the first time specifications are requested and missing + if not hasattr(self, "_specifications"): + self.setup_metadata() + return self._specifications + + @specifications.setter + def specifications( + self, specifications: Union[Specifications, Tuple[Specifications]] + ): + self._specifications = specifications + + @property + def has_setup_metadata(self): + return getattr(self, "_has_setup_metadata", False) + + @has_setup_metadata.setter + def has_setup_metadata(self, value: bool): + self._has_setup_metadata = value + + def setup_metadata(self): """Called at the beginning of training at the very beginning of Model.setup(stage="fit") Notes @@ -241,7 +299,10 @@ def setup(self, stage: Optional[str] = None): If `specifications` attribute has not been set in `__init__`, `setup` is your last chance to set it. """ - pass + + if not self.has_setup_metadata: + self.setup() + self.has_setup_metadata = True def setup_loss_func(self): pass @@ -270,18 +331,6 @@ def train_dataloader(self) -> DataLoader: collate_fn=partial(self.collate_fn, stage="train"), ) - @cached_property - def logging_prefix(self): - - prefix = f"{self.__class__.__name__}-" - if hasattr(self.protocol, "name"): - # "." has a special meaning for pytorch-lightning checkpointing - # so we remove dots from protocol names - name_without_dots = "".join(self.protocol.name.split(".")) - prefix += f"{name_without_dots}-" - - return prefix - def default_loss( self, specifications: Specifications, target, prediction, weight=None ) -> torch.Tensor: @@ -314,7 +363,7 @@ def default_loss( ]: return binary_cross_entropy(prediction, target, weight=weight) - elif specifications.problem == Problem.MONO_LABEL_CLASSIFICATION: + elif specifications.problem in [Problem.MONO_LABEL_CLASSIFICATION]: return nll_loss(prediction, target, weight=weight) else: @@ -345,6 +394,11 @@ def common_step(self, batch, batch_idx: int, stage: Literal["train", "val"]): {"loss": loss} """ + if isinstance(self.specifications, tuple): + raise NotImplementedError( + "Default training/validation step is not implemented for multi-task." + ) + # forward pass y_pred = self.model(batch["X"]) @@ -370,12 +424,17 @@ def common_step(self, batch, batch_idx: int, stage: Literal["train", "val"]): # compute loss loss = self.default_loss(self.specifications, y, y_pred, weight=weight) + + # skip batch if something went wrong for some reason + if torch.isnan(loss): + return None + self.model.log( - f"{self.logging_prefix}{stage.capitalize()}Loss", + f"loss/{stage}", loss, on_step=False, on_epoch=True, - prog_bar=True, + prog_bar=False, logger=True, ) return {"loss": loss} @@ -413,9 +472,6 @@ def val_dataloader(self) -> Optional[DataLoader]: def validation_step(self, batch, batch_idx: int): return self.common_step(batch, batch_idx, "val") - def validation_epoch_end(self, outputs): - pass - def default_metric(self) -> Union[Metric, Sequence[Metric], Dict[str, Metric]]: """Default validation metric""" msg = f"Missing '{self.__class__.__name__}.default_metric' method." @@ -426,7 +482,7 @@ def metric(self) -> MetricCollection: if self._metric is None: self._metric = self.default_metric() - return MetricCollection(self._metric, prefix=self.logging_prefix) + return MetricCollection(self._metric) def setup_validation_metric(self): metric = self.metric diff --git a/pyannote/audio/interactive/common/commands.png b/pyannote/audio/interactive/common/commands.png deleted file mode 100755 index 29517df5a..000000000 Binary files a/pyannote/audio/interactive/common/commands.png and /dev/null differ diff --git a/pyannote/audio/interactive/common/controller.js b/pyannote/audio/interactive/common/controller.js deleted file mode 100644 index cbffc2899..000000000 --- a/pyannote/audio/interactive/common/controller.js +++ /dev/null @@ -1,367 +0,0 @@ -// The MIT License (MIT) -// -// Copyright (c) 2021 CNRS -// -// Permission is hereby granted, free of charge, to any person obtaining a copy -// of this software and associated documentation files (the "Software"), to deal -// in the Software without restriction, including without limitation the rights -// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -// copies of the Software, and to permit persons to whom the Software is -// furnished to do so, subject to the following conditions: -// -// The above copyright notice and this permission notice shall be included in -// all copies or substantial portions of the Software. -// -// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -// SOFTWARE. - -// Can't create constant because prodigy reload this file every batch -var currentRegion = 0; -var regions = null; -var ids = null; -var refresh = true; - -var left = 'ArrowLeft'; -var right = 'ArrowRight'; -var startR = 'Shift'; -var endR = 'Control'; - -var PRECISION = (prodigy.config.precision / 1000); -var BEEP = prodigy.config.beep; -var EXCERPT = 1; - -var keysMap = {}; - -/** -* Handle web audio for beep -* @see beep() -*/ -var audioCtx = new(window.AudioContext || window.webkitAudioContext)(); - -/** -* Makes sure that the document is loaded before executing waitForElement() -* Add a small timeout if wavesurfer is already defined (useful when a new batch is coming), -* the time for prodigy to update wavesurfer. -* @see waitForElement() -*/ -if(document.readyState !== 'loading') { - if(typeof window.wavesurfer !== "undefined"){ - setTimeout(waitForElement,25); - }else{ - waitForElement(); - } -} else { - document.addEventListener('DOMContentLoaded', function () { - waitForElement(); - }); -} - -/** -* Compare if a region is before or after another one -* Useful to sort global variable 'ids' -* @see reloadWave() -* @param {region1} Region object -* @param {region2} Region Object -* @return {number} -1 if region1 start before, 1 otherwise, 0 if the identical -*/ -function compare(region1, region2){ - if(region1.start < region2.start){ - return -1; - }else if (region1.start > region2.start){ - return 1; - }else{ - return 0; - } -} - -/** -* Simulate a click on prodigy's label radio button -* Only useful for "review recipe" -* Note: might break with future versions of Prodigy -* @param {label} Label to click on -*/ -function clickOnLabel(label){ - document.querySelector("input[type=radio][value=\'"+label+"\']").click() -} - -/** -* Update prodigy.content with all regions from window.wavesurfer.regions.list (thus, all regions that can be seen in the interface) -* Discuss in this issue : -* https://support.prodi.gy/t/weird-interaction-between-window-prodigy-update-and-wavesurfer/5450 -*/ -function updateContent(){ - var regions = window.wavesurfer.regions.list; - var content = []; - for (var id in regions){ - var region = regions[id]; - content.push({start : region.start, end : region.end, label : region.label, id : region.id, color : region.color}); - } - window.prodigy.update({audio_spans : content}); -} - -/** -* Create a beep sound from scratch -* You can adjust the gain, frequency (here its A 440) and duration -*/ -function beep() { - if(BEEP){ - var oscillator = audioCtx.createOscillator(); - var gainNode = audioCtx.createGain(); - - oscillator.connect(gainNode); - gainNode.connect(audioCtx.destination); - - gainNode.gain.value = 0.1; - oscillator.frequency.value = 440; - oscillator.type = "square"; - - oscillator.start(); - - setTimeout( - function() { - oscillator.stop(); - }, - 150 // FIXME: should depend on the value of "precision" - ); - } -} - -/** -* Change CSS style for the selected region -* @param {e} Region object -*/ -function activeRegion(e){ - e.element.style.borderTop = "3px solid"; - e.element.style.borderBottom = "3px solid"; -} - -/** -* @see activeRegion() -* Undo CSS change -* @param {e} Region object -*/ -function deactiveRegion(e){ - e.element.style.borderTop = ""; - e.element.style.borderBottom = ""; -} - -/** -* Update global variables 'regions' and 'ids' with regions in window.wavesurfer.regions.list -* Put the first one as "active" and update the variable currentRegion -*/ -function reloadWave(){ - regions = window.wavesurfer.regions.list; - ids = Object.values(regions); - ids.sort(compare); - if(ids.length > 0){ - currentRegion = 0; - activeRegion(ids[0]); - } -} - -/** -* Switch selected region -* Update var currentRegion -* Place wavesurfer cursor at the beginning of the new region or the beginning of the file if it's a new prodigy task -* @see activeRegion() / deactiveRegion() -* @param {ids} Ids of the region to be selected -*/ -function switchCurrent(newId){ - if(ids.length > 0){ - deactiveRegion(ids[currentRegion]); - currentRegion = newId; - activeRegion(ids[newId]) - if(refresh){ - window.wavesurfer.seekTo(0); - }else{ - var time = (ids[currentRegion].start) / (window.wavesurfer.getDuration()); - window.wavesurfer.seekTo(time); - } - } -} - -/** -* Handle wavesurfer regions -* Add event listener to some wavesurfer event -*/ -function waitForElement(){ - if(typeof window.wavesurfer !== "undefined"){ - reloadWave(); - // Select created region or the first one if it's a new task - window.wavesurfer.on('region-created', function(e){ - setTimeout(function(){ - if(ids.length > 0) deactiveRegion(ids[currentRegion]); - reloadWave(); - if(refresh){ - switchCurrent(0); - }else{ - switchCurrent(ids.indexOf(e)); - } - }, 5); - }); - // Change region label (by remove the old one and create a new one with proper label) - window.wavesurfer.on('region-dblclick',function(e){ - re = window.wavesurfer.addRegion({'start' : e.start,'end' : e.end}); - e.remove(); - window.wavesurfer.fireEvent('region-update-end',re); - }); - // Select region on click - window.wavesurfer.on('region-click',function(e){ - switchCurrent(ids.indexOf(e)); - }); - // Beep when region end - window.wavesurfer.on('region-out',function(e){ - beep(); - }); - // @see updateContent() - window.wavesurfer.on('region-update-end', function(e){ - updateContent(); - }); - // @see updateContent() - // Switch selected region when deleted - window.wavesurfer.on('region-removed',function(e){ - updateContent(); - if(currentRegion == (ids.length - 1)){ - var newId = 0; - }else{ - var newId = currentRegion; - } - reloadWave(); - if(ids.length > 0) switchCurrent(newId); - }); - }else{ - setTimeout(waitForElement, 250); - } -} - -// Check if it's a new prodigy task -document.addEventListener('prodigyanswer', e => { - refresh = true; -}) - -/** -* Keyboard controller -* | Key 1 | Key 2 | Command | -* | ------------- | ------------- | ------------ | -* | Arrows left/right | [W] | Move Cursor [speed up] | -* | Shift | Arrows left/right | Change start of current segment | -* | Control | Arrows left/right | Change end of current segment | -* | Arrows up/down | | Change current segment to the next/precedent one | -* | Shift | Arrows up/[down] | Create [or remove] segment | -* | Backspace | | Remove current segment | -*/ -document.querySelector('#root').onkeydown = document.querySelector('#root').onkeyup = function(e){ - e = e || event; - keysMap[e.key] = e.type == 'keydown'; - var pos = window.wavesurfer.getCurrentTime(); - var audioEnd = window.wavesurfer.getDuration(); - var region = ids[currentRegion]; - refresh = false; - - // If Left is pressed - if(keysMap[left] && !keysMap[right]){ - // If Shift is pressed - if(keysMap[startR] && !keysMap[endR]){ - // Shortens start if possible - if((region.start - PRECISION) <= 0){ - region.update({'start' : 0}); - window.wavesurfer.fireEvent('region-update-end',region); - window.wavesurfer.play(0, region.end); - }else{ - region.update({'start' : region.start - PRECISION }); - window.wavesurfer.fireEvent('region-update-end',region); - window.wavesurfer.play(region.start, region.end); - } - // If Ctrl is pressed - }else if(keysMap[endR] && !keysMap[startR]){ - var startTime = region.end - EXCERPT; - if(startTime < region.start) startTime = region.start; - // Shortens end if possible - if((region.end - PRECISION) > region.start){ - region.update({'end' : region.end - PRECISION }); - window.wavesurfer.fireEvent('region-update-end',region); - window.wavesurfer.play(startTime, region.end); - } - }else{ - // Else change cursor position - // Speed up naviguation if W is pressed - if(keysMap['w']){ - var time = (pos - PRECISION*2) / audioEnd; - }else{ - var time = (pos - PRECISION) / audioEnd; - } - if(time < 0) time = 0; - window.wavesurfer.pause(); - window.wavesurfer.seekTo(time); - } - // If Right is pressed - }else if(keysMap[right] && !keysMap[left]){ - // If Shift is pressed - if(keysMap[startR] && !keysMap[endR]){ - // Extend start if possible - if(region.start + PRECISION < region.end){ - region.update({'start' : region.start + PRECISION }); - window.wavesurfer.fireEvent('region-update-end',region); - window.wavesurfer.play(region.start, region.end); - } - // If Ctrl is pressed - }else if(keysMap[endR] && !keysMap[startR]){ - // Extend end if possible (while keep playing the audio) - if(!window.wavesurfer.isPlaying()){ - var startTime = region.end - EXCERPT; - if(startTime < region.start) startTime = region.start; - }else{ - var startTime = pos; - } - if((region.end + PRECISION) >= audioEnd){ - region.update({'end' : audioEnd }); - window.wavesurfer.fireEvent('region-update-end',region); - }else{ - region.update({'end' : region.end + PRECISION }); - window.wavesurfer.fireEvent('region-update-end',region); - } - window.wavesurfer.play(startTime, region.end); - }else{ - // Else change cursor position - // Speed up naviguation if W is pressed - if(keysMap['w']){ - var time = (pos + PRECISION*2) / audioEnd; - }else{ - var time = (pos + PRECISION) / audioEnd; - } - if(time > 1) time = 1; - window.wavesurfer.pause(); - window.wavesurfer.seekTo(time); - } - // If Up and shift is pressed : new region - }else if (keysMap['ArrowUp'] && keysMap['Shift']){ - var fin = pos + 1; - if(fin > audioEnd) fin = audioEnd; - re = window.wavesurfer.addRegion({'start' : pos,'end' : fin}); - window.wavesurfer.fireEvent('region-update-end',re); - // If Down and Shift or Backspace: delete region - // Check backspace for diarization text field - }else if(keysMap['Backspace'] || (keysMap['ArrowDown'] && keysMap['Shift'])){ - ids[currentRegion].remove(); - // If Up/Down @see switchCurrent - }else if(keysMap['ArrowUp']){ - if(currentRegion == (ids.length - 1)){ - switchCurrent(0); - }else{ - switchCurrent(currentRegion + 1); - } - }else if(keysMap['ArrowDown']){ - if(currentRegion == 0){ - switchCurrent(ids.length - 1); - }else{ - switchCurrent(currentRegion - 1); - } - }else if(keysMap['u']){ - reloadWave(); - } -} diff --git a/pyannote/audio/interactive/common/instructions.html b/pyannote/audio/interactive/common/instructions.html deleted file mode 100644 index 0b75a98c1..000000000 --- a/pyannote/audio/interactive/common/instructions.html +++ /dev/null @@ -1,69 +0,0 @@ - - - -
- | Key 1 | -Key 2 | -Command | -
|---|---|---|
| Arrows left/right | -[W] | -Move Cursor [speed up] | -
| Shift | -Arrow left/right | -Change start of current segment | -
| Control | -Arrow left/right | -Change end of current segment | -
| Arrows up/down | -- | Change current segment to the next/precedent one | -
| Shift | -Arrow up/[down] | -Create or [remove] segment | -
| Backspace | -- | Remove current segment | -
| Spacebar | -- | Play/pause audio | -
| U | -- | Undo all actions | -
| Escape | -- | Ignore this sample | -
| Enter | -- | Validate annotation | -
![\"Open](\"https://colab.research.google.com/assets/colab-badge.svg\"){kind=icon}
"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "nxNf1l8Ye_U9"
- },
- "source": [
- "[`pyannote.audio`](https://github.com/pyannote/pyannote-audio) is an open-source toolkit written in Python for **speaker diarization**. \n",
- "\n",
- "Based on [`PyTorch`](https://pytorch.org) machine learning framework, it provides a set of trainable end-to-end neural building blocks that can be combined and jointly optimized to build speaker diarization pipelines. \n",
- "\n",
- "`pyannote.audio` also comes with pretrained [models](https://huggingface.co/models?other=pyannote-audio-model) and [pipelines](https://huggingface.co/models?other=pyannote-audio-pipeline) covering a wide range of domains for voice activity detection, speaker segmentation, overlapped speech detection, speaker embedding reaching state-of-the-art performance for most of them. \n",
- "\n",
- "**This notebook will teach you how to apply those pretrained pipelines on your own data.**\n",
- "\n",
- "Make sure you run it using a GPU (or it might otherwise be slow...)"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "42nBTa_QgooG"
- },
- "source": [
- "## Installation"
- ]
- },
- {
- "cell_type": "code",
- "metadata": {
- "id": "vJGyKTQJqdzq",
- "outputId": "28eabcb6-b64b-4e78-9dc0-a181797c3e84",
- "colab": {
- "base_uri": "https://localhost:8080/"
- }
- },
- "source": [
- "# for speechbrain\n",
- "!pip install -qq torch==1.11.0 torchvision==0.12.0 torchaudio==0.11.0 torchtext==0.12.0\n",
- "!pip install -qq speechbrain==0.5.12\n",
- "\n",
- "# pyannote.audio\n",
- "!pip install -qq pyannote.audio\n",
- "\n",
- "# for visualization purposes\n",
- "!pip install -qq moviepy ipython==7.34.0"
- ],
- "execution_count": 1,
- "outputs": [
- {
- "output_type": "stream",
- "name": "stdout",
- "text": [
- "\u001b[K |████████████████████████████████| 750.6 MB 12 kB/s \n",
- "\u001b[K |████████████████████████████████| 21.0 MB 1.4 MB/s \n",
- "\u001b[K |████████████████████████████████| 2.9 MB 49.5 MB/s \n",
- "\u001b[K |████████████████████████████████| 10.4 MB 56.5 MB/s \n",
- "\u001b[K |████████████████████████████████| 496 kB 34.8 MB/s \n",
- "\u001b[K |████████████████████████████████| 101 kB 13.2 MB/s \n",
- "\u001b[K |████████████████████████████████| 1.2 MB 58.4 MB/s \n",
- "\u001b[K |████████████████████████████████| 596 kB 56.7 MB/s \n",
- "\u001b[K |████████████████████████████████| 109 kB 75.9 MB/s \n",
- "\u001b[K |████████████████████████████████| 546 kB 74.2 MB/s \n",
- "\u001b[?25h Building wheel for hyperpyyaml (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
- "\u001b[K |████████████████████████████████| 385 kB 30.0 MB/s \n",
- "\u001b[K |████████████████████████████████| 110 kB 66.8 MB/s \n",
- "\u001b[K |████████████████████████████████| 79 kB 9.1 MB/s \n",
- "\u001b[K |████████████████████████████████| 47 kB 5.7 MB/s \n",
- "\u001b[K |████████████████████████████████| 41 kB 537 kB/s \n",
- "\u001b[K |████████████████████████████████| 129 kB 63.0 MB/s \n",
- "\u001b[K |████████████████████████████████| 419 kB 65.9 MB/s \n",
- "\u001b[K |████████████████████████████████| 51 kB 201 kB/s \n",
- "\u001b[K |████████████████████████████████| 65 kB 5.4 MB/s \n",
- "\u001b[K |████████████████████████████████| 585 kB 71.6 MB/s \n",
- "\u001b[K |████████████████████████████████| 117 kB 74.6 MB/s \n",
- "\u001b[K |████████████████████████████████| 130 kB 64.6 MB/s \n",
- "\u001b[K |████████████████████████████████| 308 kB 77.2 MB/s \n",
- "\u001b[K |████████████████████████████████| 81 kB 11.4 MB/s \n",
- "\u001b[K |████████████████████████████████| 209 kB 75.5 MB/s \n",
- "\u001b[K |████████████████████████████████| 140 kB 74.8 MB/s \n",
- "\u001b[K |████████████████████████████████| 1.1 MB 64.1 MB/s \n",
- "\u001b[K |████████████████████████████████| 58 kB 7.2 MB/s \n",
- "\u001b[K |████████████████████████████████| 144 kB 68.4 MB/s \n",
- "\u001b[K |████████████████████████████████| 271 kB 68.2 MB/s \n",
- "\u001b[K |████████████████████████████████| 94 kB 4.3 MB/s \n",
- "\u001b[K |████████████████████████████████| 78 kB 8.5 MB/s \n",
- "\u001b[K |████████████████████████████████| 112 kB 68.5 MB/s \n",
- "\u001b[K |████████████████████████████████| 147 kB 73.3 MB/s \n",
- "\u001b[K |████████████████████████████████| 49 kB 6.3 MB/s \n",
- "\u001b[?25h Building wheel for antlr4-python3-runtime (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
- " Building wheel for julius (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
- " Building wheel for pyperclip (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
- "\u001b[K |████████████████████████████████| 793 kB 22.7 MB/s \n",
- "\u001b[K |████████████████████████████████| 381 kB 72.2 MB/s \n",
- "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
- "jupyter-console 5.2.0 requires prompt-toolkit<2.0.0,>=1.0.0, but you have prompt-toolkit 3.0.30 which is incompatible.\n",
- "google-colab 1.0.0 requires ipython~=5.5.0, but you have ipython 7.34.0 which is incompatible.\u001b[0m\n",
- "\u001b[?25h"
- ]
- }
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "5pSYhjbMhTwT"
- },
- "source": [
- "# Visualization with `pyannote.core`\n",
- "\n",
- "For the purpose of this notebook, we will download and use an audio file coming from the [AMI corpus](http://groups.inf.ed.ac.uk/ami/corpus/), which contains a conversation between 4 people in a meeting room."
- ]
- },
- {
- "cell_type": "code",
- "metadata": {
- "id": "w7YuPI28qlwZ"
- },
- "source": [
- "!wget -q http://groups.inf.ed.ac.uk/ami/AMICorpusMirror/amicorpus/ES2004a/audio/ES2004a.Mix-Headset.wav\n",
- "DEMO_FILE = {'uri': 'ES2004a.Mix-Headset', 'audio': 'ES2004a.Mix-Headset.wav'}"
- ],
- "execution_count": 2,
- "outputs": []
- },
- {
- "cell_type": "markdown",
- "metadata": {
- "id": "W0_wwm3dj_bO"
- },
- "source": [
- "Because AMI is a benchmarking dataset, it comes with manual annotations (a.k.a *groundtruth*). \n",
- "Let us load and visualize the expected output of the speaker diarization pipeline.\n"
- ]
- },
- {
- "cell_type": "code",
- "metadata": {
- "id": "1pra1iVnrp5D"
- },
- "source": [
- "!wget -q https://raw.githubusercontent.com/pyannote/AMI-diarization-setup/main/only_words/rttms/test/ES2004a.rttm"
- ],
- "execution_count": 5,
- "outputs": []
- },
- {
- "cell_type": "code",
- "metadata": {
- "id": "aj1obYqhr3Qk",
- "outputId": "2b1e6b24-c355-4edb-bf8f-46c1bd7f92d2",
- "colab": {
- "base_uri": "https://localhost:8080/",
- "height": 118
- }
- },
- "source": [
- "# load groundtruth\n",
- "from pyannote.database.util import load_rttm\n",
- "_, groundtruth = load_rttm('ES2004a.rttm').popitem()\n",
- "\n",
- "# visualize groundtruth\n",
- "groundtruth"
- ],
- "execution_count": 6,
- "outputs": [
- {
- "output_type": "execute_result",
- "data": {
- "text/plain": [
- "