Spectral Denoising using 1D ResUNet

This project implements a 1D Residual U-Net (ResUNet) for denoising spectral data. The model is trained to remove noise, baseline shifts, and other artifacts from raw spectra, producing a clean version that is more suitable for analysis.

About The Project

This project provides a complete workflow for training and evaluating a 1D ResUNet model for spectral denoising. The model architecture is based on the U-Net design with residual connections, which helps in training deeper networks and achieving better performance.

The key components of this project are:

• scripts/train_resunet.py: The main script for training the denoising model.
• scripts/evaluate_model.py: A script to evaluate the performance of the trained model.
• notebooks/demo_analysis.ipynb: A Jupyter notebook demonstrating how to use the trained model for denoising a sample spectrum.

Getting Started

To get a local copy up and running, follow these simple steps.

Prerequisites

This project uses conda for environment management. Make sure you have Anaconda or Miniconda installed.

Installation

1. Clone the repo

   git clone https://github.com/nabhya8013/spectral_denoise.git
   cd spectral_denoise

2. Create and activate the Conda environment

   conda create -n spectral_env python=3.10
   conda activate spectral_env

3. Install the required packages

   pip install -r requirements.txt

Usage

To use the pretrained model for denoising your own spectral data, you can adapt the notebooks/demo_analysis.ipynb notebook. The basic steps are:

1. Load the trained model.
2. Load your raw spectral data.
3. Preprocess the data (e.g., resampling to the target length of 1024).
4. Pass the data through the model to get the denoised spectrum.

Here's a code snippet from the demo notebook:

import torch
import numpy as np
from scipy.signal import resample
from scripts.train_resunet import ResUNet1D

# Load the model
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = ResUNet1D().to(device)
model.load_state_dict(torch.load("models/resunet1d.pth", map_location=device))
model.eval()

# Load and preprocess your data
raw_spectrum = np.loadtxt("path/to/your/spectrum.txt")
raw_spectrum_resampled = resample(raw_spectrum, 1024)
noisy_tensor = torch.tensor(raw_spectrum_resampled, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device)

# Denoise
with torch.no_grad():
    denoised_tensor = model(noisy_tensor)

denoised_spectrum = denoised_tensor.cpu().squeeze().numpy()

Training

To train the model from scratch, you can run the train_resunet.py script. Make sure your data is in the data/pairs directory, with _clean.npy and _noisy.npy file pairs.

python scripts/train_resunet.py

The script will:

• Split the data into training and validation sets.
• Augment the training data by adding noise, baseline shifts, and spikes.
• Train the ResUNet1D model using a HybridLoss function.
• Save the trained model to models/resunet1d.pth.
• Evaluate the model and save the metrics to results/eval_metrics.json.

Evaluation

To evaluate the model on the validation set, you can either run the training script (which includes evaluation at the end) or run the dedicated evaluation script:

python scripts/evaluate_model.py

This will load the trained model and compute the following metrics on the validation set:

• Mean Squared Error (MSE)
• Peak Signal-to-Noise Ratio (PSNR)
• Structural Similarity Index (SSIM)
• Pearson Correlation

The results are saved in results/eval_metrics.json. The latest evaluation results are:

Metric	Value
Mean MSE	0.0019
Mean PSNR	42.68 dB
Mean SSIM	0.9927
Mean Corr	0.9911
Overall Quality	95.17%

Demonstration

The following plot, generated by notebooks/demo_analysis.ipynb, shows a comparison between the original raw spectrum, the baseline-corrected (but flawed) target, and the model's denoised output.