NOIR: Neural Operator mapping for Implicit Representations

Sidaty El Hadramy¹  ·  Nazim Haouchine²  ·  Michael Wehrli¹  ·  Philippe C. Cattin¹

¹ Department of Biomedical Engineering, University of Basel, Allschwil, Switzerland  |  ² Harvard Medical School, Brigham and Women's Hospital, Boston, United States

---

---

Overview

This repository contains the implementation of NOIR, a framework that reformulates medical imaging tasks as a function mapping problem by learning mappings between Implicit Neural Representations (INRs) using a Neural Operator. We provide an example on the Shenzhen chest X-ray dataset, covering training of the input INR, output INR, and Neural Operator, as well as a test script to reproduce our segmentation results.

---

Visualization

Running the test script produces the following sample outputs:

---

Dataset

We use the Shenzhen Hospital Chest X-ray Dataset, publicly available from the U.S. National Library of Medicine. Download it online

Once downloaded, set the dataset path in the relevant config files under config/shenzhen/ (look for the data.path field). We provide the train/val/test split used in our experiments as a JSON file in the same directory, add it within your daset folder, next to img/ and mask/ folders.

---

Repository Structure

NOIR/
├── noir/                        # Core source code
│   ├── datasets.py              # Dataset implementations (Shenzhen single, segmentation, paired)
│   ├── inrs.py                  # ModulatedSIREN implementation
│   ├── metalearning.py          # Inner/outer loop optimization of INRs (MAML-style)
│   ├── no.py                    # Neural Operator architecture
│   └── utils.py                 # Utility functions
│
├── train/                       # Training scripts
│   ├── train_input_inr.py       # Train the input INR
│   ├── train_output_inr.py      # Train the output INR
│   └── train_no.py              # Train the Neural Operator
│
├── test/                        # Test scripts
│   └── test_no.py               # End-to-end test of the Neural Operator
│
├── config/shenzhen/             # Configuration files
│   ├── train_input_inr.yaml
│   ├── train_output_inr.yaml
│   ├── train_no.yaml
│   └── test_no.yaml
│
├── models/                      # Pretrained weights
│   ├── input_inr.pth            # Input INR weights
│   ├── output_inr.pth           # Output INR weights
│   └── neural_operator.pth      # Neural Operator weights
│
├── output/                      # All results are saved here (auto-created)
│   ├── train/
│   │   ├── input_inr/           # Input INR training results
│   │   ├── output_inr/          # Output INR training results
│   │   └── no/                  # Neural Operator training results
│   └── test/
│       └── no/                  # Neural Operator test results and visualizations
│
├── assets/img/                  # Figures
├── environment.yml              # Conda environment file
└── README.md

Source code summary (noir/)

File	Description
inrs.py	Implementation of the modulated SIREN used as the INR backbone, including SirenLayer, Siren, and ModulatedSiren
metalearning.py	Handles INR optimization via inner/outer loop meta-learning, including inner_loop and outer_step
no.py	Neural Operator MLP that maps input latent codes $z_\text{in}$ to output latent codes $z_\text{out}$, with residual connections
datasets.py	Dataset classes for the Shenzhen example: ShenzhenSingleImageDataset, ShenzhenSegmentationDataset, and ShenzhenPairedDataset
utils.py	Shared utilities (resolution degradation, logging, etc.)

---

Environment Setup

We provide a conda environment file for easy setup.

1. Create the environment:

conda env create -f noir_environment.yml

2. Activate the environment:

conda activate noir

Note on CUDA: The environment.yml targets CUDA 11.8 by default. If you are on CUDA 12, edit the file and replace pytorch-cuda=11.8 with pytorch-cuda=12.1 before creating the environment.

---

Pretrained Models

We provide pretrained weights for all three components in the models/ directory:

Model	File	Description
Input INR	models/shenzhen_input_inr.pth	Modulated SIREN fitted on chest X-ray images
Output INR	models/shenzhen_output_inr.pth	Modulated SIREN fitted on lung segmentation masks
Neural Operator	models/shenzhen_neural_operator.pth	NO mapping $z_\text{in} \rightarrow z_\text{out}$

---

Testing

Download the Shenzhen dataset, set its path in config/shenzhen/test_no.yaml, and run:

python -m test.test_no --config ./config/shenzhen/test_no.yaml

The split JSON file used in our experiments is provided alongside the config. Results and visualizations will be saved to output/test/no/.

---

Training

Each component can be trained independently using its respective script and config file.

Train the Input INR (fits modulated SIRENs on input images):

python -m train.train_input_inr --config ./config/shenzhen/train_input_inr.yaml

Train the Output INR (fits modulated SIRENs on segmentation masks):

python -m train.train_output_inr --config ./config/shenzhen/train_output_inr.yaml

Train the Neural Operator (learns the latent-to-latent mapping):

python -m train.train_no --config ./config/shenzhen/train_no.yaml

Training the INRs first is required before training the Neural Operator, as the NO is trained on the latent codes produced by the fitted INRs.

---