From a919d3fcf0b97d4fd2cde4bcac5ea300df18620a Mon Sep 17 00:00:00 2001 From: yaswanth169 Date: Sat, 14 Feb 2026 00:33:50 +0530 Subject: [PATCH] feat: add auralization pipeline for inference at arbitrary positions Add inference and auralization module that enables rendering Room Impulse Responses (RIRs) at arbitrary source-listener positions using trained AV-DAR models, with optional audio convolution for spatial auralization. New files: - auralize.py: CLI entry point (alongside train.py / evaluate.py) - avdar/auralization/engine.py: InferenceEngine wrapping RirRenderer - avdar/auralization/audio.py: audio I/O and convolution utilities - avdar/auralization/README.md: feature documentation No existing files modified. --- auralize.py | 186 +++++++++++++++++++++++++++++++++ avdar/auralization/README.md | 155 +++++++++++++++++++++++++++ avdar/auralization/__init__.py | 2 + avdar/auralization/audio.py | 38 +++++++ avdar/auralization/engine.py | 108 +++++++++++++++++++ 5 files changed, 489 insertions(+) create mode 100644 auralize.py create mode 100644 avdar/auralization/README.md create mode 100644 avdar/auralization/__init__.py create mode 100644 avdar/auralization/audio.py create mode 100644 avdar/auralization/engine.py diff --git a/auralize.py b/auralize.py new file mode 100644 index 0000000..06f26b9 --- /dev/null +++ b/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, + ) diff --git a/avdar/auralization/README.md b/avdar/auralization/README.md new file mode 100644 index 0000000..038745e --- /dev/null +++ b/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. diff --git a/avdar/auralization/__init__.py b/avdar/auralization/__init__.py new file mode 100644 index 0000000..e903510 --- /dev/null +++ b/avdar/auralization/__init__.py @@ -0,0 +1,2 @@ +from ..utils.import_utils import import_children +import_children(__file__, __name__) diff --git a/avdar/auralization/audio.py b/avdar/auralization/audio.py new file mode 100644 index 0000000..a228ada --- /dev/null +++ b/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) diff --git a/avdar/auralization/engine.py b/avdar/auralization/engine.py new file mode 100644 index 0000000..b399baa --- /dev/null +++ b/avdar/auralization/engine.py @@ -0,0 +1,108 @@ +import torch +import numpy as np +import logging + +from scipy.spatial.transform import Rotation + +from ..model.renderer import RirRenderer +from ..geometry.pathspace import SpecularPathSampler +from .audio import convolve_rir, normalize_audio + +logger = logging.getLogger(__name__) + + +class InferenceEngine: + """Wraps a trained RirRenderer for inference at arbitrary positions.""" + + def __init__(self, config, renderer, path_sampler, dataset, device): + self.config = config + self.renderer = renderer + self.path_sampler = path_sampler + self.dataset = dataset + self.device = device + + self.renderer.eval() + + self.sample_rate = config.dataset.sample_rate + self.speed_of_sound = config.dataset.options.speed_of_sound + + @torch.no_grad() + def render_rir(self, source_xyz, listener_xyz, source_orientation=None): + """Render RIR at a source-listener position. Returns (rir, metadata).""" + source_xyz_np = np.array(source_xyz, dtype=np.float32) + listener_xyz_np = np.array(listener_xyz, dtype=np.float32) + + if source_orientation is None: + source_quat = np.array([0, 0, 0, 1], dtype=np.float32) + else: + source_quat = np.array(source_orientation, dtype=np.float32) + + # Rotation matrix from quaternion + rotation = None + try: + rot_mat = Rotation.from_quat(source_quat).as_matrix() + rotation = torch.from_numpy(rot_mat.astype(np.float32)).to(self.device) + except Exception: + pass + + # Beam tracing + mc_samples = self.path_sampler.fast_sample(source_xyz_np, listener_xyz_np) + + source_t = torch.tensor(source_xyz_np, dtype=torch.float32).to(self.device) + listener_t = torch.tensor(listener_xyz_np, dtype=torch.float32).to(self.device) + quat_t = torch.tensor(source_quat, dtype=torch.float32).to(self.device) + + # Forward pass — same signature as eval_step_rir in run.py + pred_dict = self.renderer( + None, None, None, None, + rotation, source_t, listener_t, quat_t, + mc_samples=mc_samples, + ) + + rir = pred_dict['rir_full'].detach().cpu().numpy() + + metadata = { + 'source_xyz': source_xyz_np.tolist(), + 'listener_xyz': listener_xyz_np.tolist(), + 'source_orientation': source_quat.tolist(), + 'sample_rate': self.sample_rate, + 'rir_length_samples': len(rir), + 'rir_duration_seconds': float(len(rir) / self.sample_rate), + } + + logger.info(f"Rendered RIR: {len(rir)} samples ({len(rir)/self.sample_rate:.3f}s)") + return rir, metadata + + @torch.no_grad() + def auralize(self, source_xyz, listener_xyz, input_audio, sr, + source_orientation=None): + """Convolve input audio with rendered RIR. Returns (auralized, rir, metadata).""" + if sr != self.sample_rate: + logger.warning( + f"Input sr ({sr}) != model sr ({self.sample_rate}). " + f"Ensure sample rates match for correct results." + ) + + rir, metadata = self.render_rir(source_xyz, listener_xyz, source_orientation) + + auralized = convolve_rir(input_audio, rir) + auralized = normalize_audio(auralized) + + metadata['input_audio_samples'] = len(input_audio) + metadata['output_audio_samples'] = len(auralized) + + return auralized, rir, metadata + + @torch.no_grad() + def batch_render(self, positions): + """Render RIRs for multiple source-listener pairs.""" + results = [] + for i, pos in enumerate(positions): + logger.info(f"Rendering {i+1}/{len(positions)}") + rir, metadata = self.render_rir( + source_xyz=pos['source_xyz'], + listener_xyz=pos['listener_xyz'], + source_orientation=pos.get('source_orientation', None), + ) + results.append({'rir': rir, 'metadata': metadata}) + return results