Pad the input so the result is not shortened by the receptive field

microtcn_runtime.py

@@ -177,6 +177,33 @@ def _load_model() -> MicroTCN:
     return _MODEL
 
 
+_RECEPTIVE_FIELD = None
+
+
+def _receptive_field(model, device) -> int:
+    """
+    How many samples the network consumes, measured once from the model itself.
+
+    Args:
+        model: The loaded network.
+        device: Where to run the measurement.
+
+    Returns:
+        deficit (int): Samples the output is short of the input, 0 if none.
+    """
+
+    global _RECEPTIVE_FIELD
+
+    if _RECEPTIVE_FIELD is None:
+        probe = torch.zeros(1, 1, SAMPLE_RATE, device=device)
+        params = torch.zeros(1, 1, 2, device=device)
+        with torch.inference_mode():
+            out = model(probe, params)
+        _RECEPTIVE_FIELD = max(0, probe.shape[-1] - out.shape[-1])
+
+    return _RECEPTIVE_FIELD
+
+
 def process_audio(
     audio_path: str,
     limit_mode: float,
@@ -191,7 +218,16 @@ def process_audio(
     if sr != SAMPLE_RATE:
         waveform = torchaudio.functional.resample(waveform, sr, SAMPLE_RATE)
 
+    # The network does not pad, so it returns its receptive field fewer samples
+    # than it is given. Left-padding by that deficit, with the first sample
+    # repeated so the pad is not a step into silence, gives a result the same
+    # length as the input and aligned with it.
+    pad = _receptive_field(model, device)
     x = waveform.unsqueeze(0).to(device)
+
+    if pad > 0:
+        x = torch.cat([x[..., :1].expand(-1, -1, pad), x], dim=-1)
+
     p = torch.tensor(
         [[[limit_mode, peak_reduction]]], dtype=torch.float32, device=device
     )
@@ -200,4 +236,10 @@ def process_audio(
         y = model(x, p)
 
     output = y.squeeze().cpu().numpy().astype(np.float32)
+
+    # Guard against a deficit that is not exactly the measured one
+    wanted = waveform.shape[-1]
+    if output.shape[-1] > wanted:
+        output = output[..., :wanted]
+
     return output, SAMPLE_RATE