# CSC2529 Rain Removal Diffusion Pipeline

This repository implements a diffusion-based rain removal system with optional attention-guided conditioning. It supports both a light-weight `SimpleUNet` backbone and a DDPM-style `DiffusionUNet`, and it can train/evaluate multiple configuration scenarios automatically.

## Core Components

- **config.yaml** – Central configuration for data, model backbones, attention toggle, diffusion hyperparameters, and training settings.
- **attention_net.py** – Encoder–decoder network that predicts a single-channel rain mask; used to augment the conditioning tensor when attention is enabled.
- **simple_unet.py** – Compact UNet that operates on a concatenation of diffuse (`x_t`) and conditioning channels. Inputs are parameterised by `x_ch` (diffusion) and `cond_ch` (conditioning) and checked at runtime.
- **ddpm_unet.py** – Full DDPM UNet with residual and attention blocks. Validates diffusion/conditioning channel counts and supports the attention mask via an extra conditioning channel.
- **model.py** – Builds the `RainDiffusionModel`, wires the AttentionNet when requested, and exposes the diffusion sampling/prediction utilities. Handles channel bookkeeping for both backbones and optional attention warm-start.
- **train.py** – Training loop with EMA tracking, evaluation checkpoints, PSNR/SSIM logging, and best-checkpoint selection based on reconstruction loss during evaluation intervals.
- **evaluate.py** – Evaluation entry point that loads a trained model and reports quantitative metrics on the configured test split.
- **pipeline.py** – Scenario runner that generates temporary configs and executes the four main experiments (simple/ddpm × attention on/off) with lightweight preview diagnostics.
- **dataloader.py** – Dataset and dataloader utilities for the RainDS dataset variants.
- **utils.py** – Helper routines for configuration loading, metrics, EMA maintenance, and shared utilities.

## Configuration Overview

`config.yaml` groups settings beneath `data`, `model`, `training`, and `evaluation`:

- **data** – Paths, image size, variant selection, and dataloader threading.
- **model** – Backbone choice, attention toggle (`model.attention.enabled`), channel splits for both backbones, diffusion schedule, and EMA defaults.
- **training** – Optimisation controls plus lightweight preview evaluation knobs (`preview_eval_interval`, `preview_eval_samples`, `preview_save_dir`, `preview_save_images`) that drive quick sanity checks during training.
- **evaluation** – Settings for the full test-time sweep (output directories, worker count, optional sample saving).

Changing the backbone or enabling the mask requires no code edits; channel counts are adjusted automatically when `model.attention.enabled` is `true`.

## Model Construction (`model.py`)

`build_model(cfg, device)` executes the following steps:

1. Reads the backbone (`simple_unet` or `ddpm_unet`).
2. Checks `model.attention.enabled`. When enabled it:
   - Instantiates `AttentionNet` with configurable channels.
   - Optionally loads weights from `model.attention.pretrained_path`.
   - Adds one conditioning channel to the selected backbone.
3. Creates the backbone (`SimpleUNet` or `DiffusionUNet`) with consistent diffusion/conditioning channel totals.
4. Constructs `RainDiffusionModel`, which stores the attention module (if any) and handles noise prediction, diffusion sampling, and denoising.

Inside `RainDiffusionModel.predict_eps`, the attention mask is concatenated to the conditioning tensor before invoking the backbone. When attention is off, the backbone runs with its standard inputs.

## Model Architecture (DDPM UNet + Attention)

When `model.backbone` is set to `ddpm_unet` and `model.attention.enabled` is `true`, the data flow is:

- **Conditioning Stack** – The rainy input `y` is passed directly to the UNet and also through `AttentionNet`, producing a sigmoid mask `m ∈ [0, 1]` (shape `B×1×H×W`). The concatenated tensor `[y, m]` forms the conditioning channels consumed by `DiffusionUNet`.
- **Diffusion Input** – The noisy sample `x_t` is kept separate (`in_ch=3`). `DiffusionUNet` validates that `x_t` carries exactly `self.x_ch` channels and concatenates `x_t` with the conditioning stack, yielding the full input tensor.
- **Time & Variant Embeddings** – `get_timestep_embedding` projects the diffusion timestep into a sinusoidal vector, which is then pushed through an MLP to produce `temb`. If variant indices are supplied (rain streak, drop, combined), they are embedded with `nn.Embedding` and added to `temb`, allowing each rain type to bias the network features.
- **Down/Up Sampling** – The network follows the DDPM reference architecture: residual blocks with GroupNorm + Swish, optional self-attention at user-specified resolutions, strided downsampling, and mirrored upsampling with skip connections.
- **Output Head** – After the final upsampling stage, a GroupNorm + Swish pair feeds the concluding `Conv2d`, predicting `ε̂` in RGB space.

This setup lets the model remain type-aware via variant embeddings while the learned attention mask highlights rain-heavy regions, giving the denoiser both global context (timesteps, rain type) and spatial guidance (mask).

## Training (`train.py`)

Key training features:

- **Variant Handling** – `_to_variant_idx` converts batch variant labels into indices for backbone conditioning. Missing variant entries are supported.
- **Loss Terms** – Noise prediction MSE, reconstruction L1 (`lambda_rec`), and optional pairwise consistency (`lambda_cons`).
- **Mixed Precision** – Uses `torch.cuda.amp` scaler (if CUDA is active).
- **EMA** – Maintains an exponential moving average of model parameters for evaluation and checkpointing.
- **Evaluation Loop** – Every `training.preview_eval_interval` epochs the loop:
  - Swaps model weights to EMA (when enabled) for inference.
  - Computes PSNR/SSIM over `training.preview_eval_samples` batches (or the full loader when set to `null`).
  - Optionally saves a preview comparison (`training.preview_save_images`).
  - Tracks reconstruction loss accumulated during the training epoch and saves `model_best.pt` whenever the average reconstruction loss improves.
- **Checkpoints** – Always writes `model_last.pt`, keeps `model_best.pt` (best reconstruction loss), and saves `model_ema.pt` at the end (final EMA snapshot).

## Evaluation (`evaluate.py`)

Loads the requested configuration and checkpoint, then runs inference over the entire test loader with PSNR/SSIM reporting. By default it only logs metrics; set `evaluation.save_images = true` to also dump composite visuals. Shares preprocessing with the training dataloader to ensure consistency.

## Experiments (`pipeline.py`)

Runs four predefined experiments:

1. Simple UNet, attention disabled
2. Simple UNet, attention enabled
3. DDPM UNet, attention disabled
4. DDPM UNet, attention enabled

For each scenario the pipeline:

- Duplicates `config.yaml`, applies overrides for backbone/attention and sets scenario-specific output folders.
- Forces `training.preview_eval_interval = 1` and `training.preview_eval_samples = 1` so each one-epoch smoke test still emits a quick diagnostic.
- Invokes `train` then `evaluate` sequentially.

### Running the Experiments

```bash
python pipeline.py
```

Outputs for each scenario are placed under `outputs/<scenario_name>` with evaluation samples in the corresponding subfolder.

## Manual Training & Evaluation

Train with the default configuration:

```bash
python train.py --config config.yaml
```

Evaluate a trained model:

```bash
python evaluate.py --config config.yaml --checkpoint outputs/run/model_best.pt
```

(Adjust paths to match saved checkpoints.)

## Attention Integration Notes

- When attention is enabled, the rainy input (`y`) feeds both the UNet conditioner and the `AttentionNet`.
- The attention mask output is concatenated along the channel dimension with the conditioning tensor before passing to the backbone.
- Both backbones validate channel counts to avoid silent shape mismatches during development or runtime reconfiguration.

## Dataset Layout

The repository expects the RainDS dataset under `dataset/` with the following structure:

```
dataset/
  RainDS_syn/
    train/
      raindrop/
      rainstreak/
      rainstreak_raindrop/
      gt/
    test/
      ...
  train_set/
    ...
  test_set/
    ...
```

Adjust `data.root`, `train_variant`, and `test_variant` in `config.yaml` to point to the desired split or subset.

## Checkpoint Summary

- `model_last.pt` – Most recent training weights (online network).
- `model_best.pt` – Best reconstruction-loss checkpoint captured at evaluation intervals (EMA weights when enabled).
- `model_ema.pt` – Final EMA snapshot after training completes.

## Reproducibility

- Use `training.seed` for deterministic behaviour (sets Python, NumPy, and PyTorch seeds and seeds CUDA when available).
- Device selection defaults to CUDA if available; override with `training.device`.

## Extending the Project

- Modify `config.yaml` to experiment with different diffusion schedules, channel multipliers, or attention settings.
- Add new backbones by extending `model.py` and updating the configuration schema.
- Integrate additional metrics in the evaluation loop by extending `train.py` and `evaluate.py`.

For further experimentation, clone configurations, tweak overrides, and run via `pipeline.py` or by calling `train`/`evaluate` directly.
