feat: add auralization pipeline for inference at arbitrary positions

auralize.py

@@ -0,0 +1,186 @@
+import argparse
+import json
+import logging
+import pathlib
+import sys
+
+import numpy as np
+import torch
+import hydra
+from omegaconf import OmegaConf
+
+from avdar.core.base_config import BaseConfig
+from avdar.core.io import build_from_config
+from avdar.auralization.engine import InferenceEngine
+from avdar.auralization.audio import load_audio, save_audio, normalize_audio
+from avdar.geometry.pathspace import SpecularPathSampler
+from avdar.model.renderer import RirRenderer
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description='Render RIRs and auralize audio at arbitrary positions'
+    )
+    parser.add_argument('--config_dir', type=str, required=True,
+                        help='Hydra training run directory')
+    parser.add_argument('--state_dict_name', type=str, default='weight_final.pt',
+                        help='Checkpoint filename within config_dir')
+    parser.add_argument('--device', type=str, default='cuda:0')
+
+    parser.add_argument('--source_xyz', type=float, nargs=3, default=None,
+                        metavar=('X', 'Y', 'Z'))
+    parser.add_argument('--listener_xyz', type=float, nargs=3, default=None,
+                        metavar=('X', 'Y', 'Z'))
+    parser.add_argument('--source_orientation', type=float, nargs=4, default=None,
+                        metavar=('QX', 'QY', 'QZ', 'QW'),
+                        help='Quaternion [x,y,z,w], default: identity')
+
+    parser.add_argument('--positions_file', type=str, default=None,
+                        help='JSON file with list of position dicts')
+    parser.add_argument('--input_audio', type=str, default=None,
+                        help='Dry audio file to auralize')
+    parser.add_argument('--output_dir', type=str, required=True)
+    return parser.parse_args()
+
+
+def main(config_dir, state_dict_name, device, positions, input_audio_path,
+         output_dir):
+
+    # Load config (same pattern as evaluate.py)
+    with hydra.initialize_config_dir(
+        config_dir=str(pathlib.Path(config_dir).absolute()),
+        version_base="1.2"
+    ):
+        config = hydra.compose(config_name='config', overrides=[
+            f'device={device}',
+            f'state_dict_path={pathlib.Path(config_dir) / state_dict_name}',
+            'no_terminal=True',
+        ])
+
+    cache_dir = (
+        pathlib.Path(config.working_dir)
+        / (config.dataset.name + '_' + config.dataset.scene_name)
+    )
+
+    # Build model and datasets
+    build_dict = build_from_config(
+        config, working_dir=config.working_dir,
+        cache_dir=cache_dir, resume=True, inference_only=True,
+    )
+
+    dataset = (
+        build_dict.get('dataset_inference')
+        or build_dict.get('dataset_test')
+        or build_dict.get('dataset_val')
+        or build_dict['dataset_train']
+    )
+
+    rir_renderer: RirRenderer = build_dict['rir_renderer'].to(device)
+    rir_renderer.eval()
+
+    # Path sampler
+    max_path_length = config.train.max_bounce
+    path_sampler = SpecularPathSampler.from_config(
+        config.train['sampler_opts'],
+        max_path_length,
+        dataset.get_mesh_path(),
+    )
+
+    engine = InferenceEngine(
+        config=config,
+        renderer=rir_renderer,
+        path_sampler=path_sampler,
+        dataset=dataset,
+        device=device,
+    )
+
+    output_path = pathlib.Path(output_dir)
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    # Load input audio if provided
+    input_audio = None
+    input_sr = None
+    if input_audio_path is not None:
+        input_audio, input_sr = load_audio(
+            input_audio_path, target_sr=config.dataset.sample_rate
+        )
+
+    # Render
+    all_metadata = []
+    for i, pos in enumerate(positions):
+        prefix = f"position_{i:04d}" if len(positions) > 1 else "output"
+        logger.info(
+            f"[{i+1}/{len(positions)}] "
+            f"src={pos['source_xyz']} -> lst={pos['listener_xyz']}"
+        )
+
+        if input_audio is not None:
+            auralized, rir, metadata = engine.auralize(
+                source_xyz=pos['source_xyz'],
+                listener_xyz=pos['listener_xyz'],
+                input_audio=input_audio,
+                sr=input_sr,
+                source_orientation=pos.get('source_orientation'),
+            )
+            auralized_path = output_path / f"{prefix}_auralized.wav"
+            save_audio(auralized, str(auralized_path), input_sr)
+            metadata['auralized_audio_path'] = str(auralized_path)
+        else:
+            rir, metadata = engine.render_rir(
+                source_xyz=pos['source_xyz'],
+                listener_xyz=pos['listener_xyz'],
+                source_orientation=pos.get('source_orientation'),
+            )
+
+        rir_path = output_path / f"{prefix}_rir.wav"
+        save_audio(normalize_audio(rir), str(rir_path), config.dataset.sample_rate)
+        metadata['rir_path'] = str(rir_path)
+        all_metadata.append(metadata)
+
+    # Save metadata
+    metadata_path = output_path / "metadata.json"
+    with open(metadata_path, 'w') as f:
+        json.dump(all_metadata, f, indent=2)
+
+    logger.info(f"Done. {len(positions)} RIR(s) saved to {output_path}")
+
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    has_single = args.source_xyz is not None and args.listener_xyz is not None
+    has_batch = args.positions_file is not None
+
+    if not has_single and not has_batch:
+        print("Error: specify --source_xyz + --listener_xyz, or --positions_file",
+              file=sys.stderr)
+        sys.exit(1)
+    if has_single and has_batch:
+        print("Error: use --source_xyz/--listener_xyz or --positions_file, not both",
+              file=sys.stderr)
+        sys.exit(1)
+
+    if has_single:
+        positions = [{
+            'source_xyz': args.source_xyz,
+            'listener_xyz': args.listener_xyz,
+            'source_orientation': args.source_orientation,
+        }]
+    else:
+        with open(args.positions_file, 'r') as f:
+            positions = json.load(f)
+
+    main(
+        config_dir=args.config_dir,
+        state_dict_name=args.state_dict_name,
+        device=args.device,
+        positions=positions,
+        input_audio_path=args.input_audio,
+        output_dir=args.output_dir,
+    )

avdar/auralization/README.md

@@ -0,0 +1,155 @@
+# Auralization Module
+
+This module provides inference and auralization capabilities for trained AV-DAR models. Given a trained checkpoint, it renders Room Impulse Responses (RIRs) at arbitrary source-listener positions in the scene and optionally convolves them with input audio — enabling spatial audio rendering at novel positions without retraining.
+
+## Directory Structure
+
+```
+avdar/auralization/
+├── __init__.py     # package init
+├── audio.py        # audio I/O: load, save, convolve, normalize
+├── engine.py       # InferenceEngine: wraps RirRenderer for inference
+└── README.md       # this file
+
+auralize.py             # CLI entry point (at repo root, alongside train.py / evaluate.py)
+```
+
+## Quick Start
+
+### Render a single RIR
+
+```bash
+python auralize.py \
+    --config_dir ./outputs/HAA-Classroom-16K/2025-10-20_12-00-00 \
+    --source_xyz 1.0 2.0 0.5 \
+    --listener_xyz 3.0 1.5 0.5 \
+    --output_dir ./auralized_output
+```
+
+This produces:
+- `output_rir.wav` — the rendered RIR
+- `metadata.json` — source/listener positions, sample rate, RIR length
+
+### Auralize an input audio file
+
+```bash
+python auralize.py \
+    --config_dir ./outputs/HAA-Classroom-16K/2025-10-20_12-00-00 \
+    --source_xyz 1.0 2.0 0.5 \
+    --listener_xyz 3.0 1.5 0.5 \
+    --input_audio ./dry_speech.wav \
+    --output_dir ./auralized_output
+```
+
+This additionally produces:
+- `output_auralized.wav` — the input audio convolved with the rendered RIR
+
+### Batch render multiple positions
+
+Create a `positions.json` file:
+
+```json
+[
+    {"source_xyz": [1.0, 2.0, 0.5], "listener_xyz": [3.0, 1.5, 0.5]},
+    {"source_xyz": [1.0, 2.0, 0.5], "listener_xyz": [5.0, 3.0, 0.5]},
+    {"source_xyz": [2.0, 1.0, 0.5], "listener_xyz": [4.0, 2.5, 0.5],
+     "source_orientation": [0, 0, 0.707, 0.707]}
+]
+```
+
+```bash
+python auralize.py \
+    --config_dir ./outputs/HAA-Classroom-16K/2025-10-20_12-00-00 \
+    --positions_file ./positions.json \
+    --output_dir ./auralized_output
+```
+
+Produces `position_0000_rir.wav`, `position_0001_rir.wav`, etc.
+
+## CLI Reference
+
+| Argument | Required | Default | Description |
+|---|---|---|---|
+| `--config_dir` | Yes | — | Path to a Hydra training run directory |
+| `--state_dict_name` | No | `weight_final.pt` | Checkpoint filename within `config_dir` |
+| `--device` | No | `cuda:0` | Compute device |
+| `--source_xyz X Y Z` | * | — | Source position in 3D |
+| `--listener_xyz X Y Z` | * | — | Listener position in 3D |
+| `--source_orientation QX QY QZ QW` | No | identity | Source rotation quaternion |
+| `--positions_file` | * | — | JSON with list of position dicts |
+| `--input_audio` | No | — | Dry audio file to convolve with RIR |
+| `--output_dir` | Yes | — | Output directory for results |
+
+\* Either `--source_xyz` + `--listener_xyz` or `--positions_file` is required (not both).
+
+## Python API
+
+The `InferenceEngine` class can be used directly in scripts:
+
+```python
+from avdar.auralization.engine import InferenceEngine
+from avdar.auralization.audio import load_audio, save_audio
+
+# ... build config, renderer, path_sampler, dataset as in auralize.py ...
+
+engine = InferenceEngine(config, renderer, path_sampler, dataset, device)
+
+# Render RIR
+rir, metadata = engine.render_rir(
+    source_xyz=[1.0, 2.0, 0.5],
+    listener_xyz=[3.0, 1.5, 0.5]
+)
+
+# Auralize
+audio, sr = load_audio("speech.wav", target_sr=16000)
+auralized, rir, metadata = engine.auralize(
+    source_xyz=[1.0, 2.0, 0.5],
+    listener_xyz=[3.0, 1.5, 0.5],
+    input_audio=audio, sr=sr
+)
+save_audio(auralized, "output.wav", sr)
+
+# Batch render
+results = engine.batch_render([
+    {"source_xyz": [1,2,0.5], "listener_xyz": [3,1.5,0.5]},
+    {"source_xyz": [1,2,0.5], "listener_xyz": [5,3,0.5]},
+])
+```
+
+## Output Format
+
+Each run produces:
+
+| File | Description |
+|---|---|
+| `*_rir.wav` | Rendered Room Impulse Response (peak-normalized, int16) |
+| `*_auralized.wav` | Input audio convolved with RIR (only if `--input_audio` provided) |
+| `metadata.json` | Positions, sample rate, RIR length, file paths |
+
+### metadata.json example
+
+```json
+[
+    {
+        "source_xyz": [1.0, 2.0, 0.5],
+        "listener_xyz": [3.0, 1.5, 0.5],
+        "source_orientation": [0, 0, 0, 1],
+        "sample_rate": 16000,
+        "rir_length_samples": 32000,
+        "rir_duration_seconds": 2.0,
+        "rir_path": "./auralized_output/output_rir.wav"
+    }
+]
+```
+
+## Architecture
+
+The auralization pipeline uses the same trained components as the evaluation loop:
+
+1. **Config loading** — Hydra `initialize_config_dir` + `compose` (same as `evaluate.py`)
+2. **Model building** — `build_from_config()` with `inference_only=True`
+3. **Beam tracing** — `SpecularPathSampler` traces specular reflection paths from source through the scene mesh to the listener
+4. **Neural rendering** — `RirRenderer.forward()` combines early reflections (specular MLP), diffuse field (positional encoding network), and late reverberation to produce the full RIR
+5. **Convolution** — FFT-based convolution of dry audio with the rendered RIR
+
+No modifications to any existing files are required.

avdar/auralization/__init__.py

@@ -0,0 +1,2 @@
+from ..utils.import_utils import import_children
+import_children(__file__, __name__)

avdar/auralization/audio.py

@@ -0,0 +1,38 @@
+import numpy as np
+import scipy.io.wavfile as wavfile
+import scipy.signal
+
+import librosa
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def load_audio(path, target_sr=None):
+    """Load audio file and optionally resample. Returns (audio, sr)."""
+    audio, sr = librosa.load(path, sr=target_sr, mono=True)
+    logger.info(f"Loaded {path} (sr={sr}, {len(audio)/sr:.2f}s)")
+    return audio, sr
+
+
+def save_audio(audio, path, sr):
+    """Save audio signal as int16 .wav file."""
+    audio_clipped = np.clip(audio, -1.0, 1.0)
+    audio_int16 = (audio_clipped * 32767).astype(np.int16)
+    wavfile.write(path, sr, audio_int16)
+    logger.info(f"Saved {path} (sr={sr}, {len(audio)/sr:.2f}s)")
+
+
+def convolve_rir(audio, rir):
+    """FFT-based convolution of audio with RIR, truncated to input length."""
+    convolved = scipy.signal.fftconvolve(audio, rir, mode='full')
+    return convolved[:len(audio)]
+
+
+def normalize_audio(audio, target_peak=0.95):
+    """Peak-normalize audio signal."""
+    peak = np.max(np.abs(audio))
+    if peak < 1e-8:
+        return audio
+    return audio * (target_peak / peak)