# Plug-and-Play Dehazing with CNN Priors

**CSC2529 Course Project**  
**Author**: Zherui Zhang  
**Date**: December 2025

---

## Overview

This code implements **single-image dehazing** using **Plug-and-Play (PnP) optimization** with CNN denoisers as learned priors. The project compares multiple denoising architectures (DnCNN, DRUNet, DiffPIR) with the Dark Channel Prior baseline within a unified physics-based framework.

**Purpose**: Generate comprehensive comparison results for the course project report.

---

## Project Structure

```
submission/
├── README.md                    # This file
├── requirements.txt             # Python dependencies
│
├── data/                        # Test dataset (5 images)
│   ├── indoor/
│   │   ├── hazy/               # 3 indoor hazy images
│   │   └── gt/                 # 3 indoor ground truth
│   └── outdoor/
│       ├── hazy/               # 2 outdoor hazy images
│       └── gt/                 # 2 outdoor ground truth
│
├── weights/                     # Auto-downloaded pretrained models
│   ├── dncnn_color_blind.pth   # DnCNN weights (auto-download)
│   └── drunet_color.pth        # DRUNet weights (auto-download)
│
├── pnp/                         # Core PnP framework
│   ├── optimizer.py            # PnP-ADMM optimization engine
│   ├── haze_model.py           # Atmospheric scattering model
│   ├── dncnn.py                # DnCNN denoiser loader
│   ├── modern_denoisers.py     # DRUNet wrapper
│   ├── download_utils.py       # Auto-download pretrained weights
│   ├── guided_filter.py        # Edge-preserving filtering
│   ├── networks/
│   │   └── drunet.py           # DRUNet architecture
│   └── diffusion/              # Diffusion-based experiments
│       ├── diffpir_dehaze.py   # DiffPIR implementation
│       ├── guided_diffusion_model.py
│       └── haze_projection.py
│
├── baselines/                   # Baseline methods
│   ├── dcp.py                  # Dark Channel Prior (He et al. 2009)
│   └── guided_filter.py        # Fast edge-aware filtering
│
├── utils/                       # Evaluation utilities
│   └── metrics.py              # PSNR, SSIM metrics
│
├── scripts/                     # Main experiment script
│   └── generate_final_comparison.py  # Generates all comparison results
│
└── outputs/                     # Generated outputs (created on run)
    └── final_comparison/
        ├── comprehensive_comparison.png  # Main figure
        ├── summary_statistics.csv
        └── summary_statistics.txt
```

---

## How to Run

### Prerequisites

**System Requirements**:
- Python 3.8 or higher
- CUDA-capable GPU recommended (works on CPU but slower)
- ~200 MB disk space for pretrained weights

### Step 1: Install Dependencies

```bash
# Create virtual environment (recommended)
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install required packages
pip install -r requirements.txt
```

**Key Dependencies**:
- PyTorch 1.12+
- OpenCV 4.5+
- NumPy, scikit-image, matplotlib, tqdm

### Step 2: Run the Comparison Script

```bash
python scripts/generate_final_comparison.py
```

**What happens**:
1. Auto-downloads pretrained weights from GitHub on first run (~200 MB)
   - DnCNN weights: `weights/dncnn_color_blind.pth`
   - DRUNet weights: `weights/drunet_color.pth`
2. Processes 5 test images with 8 methods:
   - Hazy (input)
   - DCP baseline
   - PnP-DnCNN (100 iterations)
   - PnP-DnCNN (200 iterations)
   - PnP-DRUNet (100 iterations)
   - PnP-DRUNet (200 iterations)
   - DiffPIR
   - Ground Truth
3. Generates comprehensive comparison figure with PSNR/SSIM metrics
4. Saves summary statistics (CSV and TXT)

**Output Location**: `outputs/final_comparison/`
- `comprehensive_comparison.png` - Main visual comparison (8×5 grid)
- `summary_statistics.csv` - Quantitative results table
- `summary_statistics.txt` - Formatted text summary

### Step 3: View Results

The main output file is `outputs/final_comparison/comprehensive_comparison.png`, which shows:
- Each row: one test image
- Each column: one method
- Metrics displayed on each result (PSNR/SSIM vs ground truth)

---

## Methods Implemented

### 1. Dark Channel Prior (DCP)
Classic physics-based method (He et al. 2009) - no learning required.

### 2. PnP-DnCNN
Plug-and-play optimization with DnCNN denoiser (Zhang et al. 2017):
- 17-layer CNN architecture
- Tested at 100 and 200 iterations
- Uses color blind denoiser

### 3. PnP-DRUNet  
Plug-and-play with modern U-Net denoiser (Zhang et al. 2021):
- Residual U-Net with skip connections
- Tested at 100 and 200 iterations
- More powerful but potentially unstable

### 4. DiffPIR
Experimental diffusion-based approach (Zhu et al. 2023):
- Integrates diffusion models with PnP framework
- Included for completeness (experimental baseline)

---

## Technical Details

### Physics-Based Haze Model

The atmospheric scattering model used:
```
I(x) = J(x) · t(x) + A · (1 - t(x))
```
where:
- **I**: Observed hazy image
- **J**: Scene radiance (clean image we want to recover)
- **t**: Transmission map (depth-dependent)
- **A**: Atmospheric light (global ambient)

### PnP-ADMM Framework

The optimization alternates between:
1. **Data fidelity step**: Enforce haze model consistency
2. **Prior step**: Apply CNN denoiser
3. **ADMM update**: Maintain variable consensus

This allows integrating pretrained denoisers without task-specific training.

---

## Troubleshooting

### Issue: "CUDA out of memory"
**Solution**: The script will automatically fall back to CPU if GPU memory is insufficient. To force CPU mode:
```bash
export CUDA_VISIBLE_DEVICES=-1
python scripts/generate_final_comparison.py
```

### Issue: "Connection timeout" during weight download
**Solution**: Pretrained weights are downloaded from GitHub. If download fails:
1. Check internet connection
2. Manually download from [KAIR releases](https://github.com/cszn/KAIR/releases/tag/v1.0)
3. Place files in `weights/` directory:
   - `dncnn_color_blind.pth`
   - `drunet_color.pth`

### Issue: "ModuleNotFoundError"
**Solution**: Ensure all dependencies are installed:
```bash
pip install -r requirements.txt
```

### Issue: Script runs but no output generated
**Solution**: Check that `data/` directory contains the test images. The script expects:
- `data/indoor/hazy/` - 3 PNG files
- `data/indoor/gt/` - 3 PNG files  
- `data/outdoor/hazy/` - 2 JPG files
- `data/outdoor/gt/` - 2 PNG files

---

## References

### Core Methods
1. **He et al.** - "Single Image Haze Removal Using Dark Channel Prior" (CVPR 2009)
2. **Zhang et al.** - "Beyond a Gaussian Denoiser: Residual Learning of Deep CNN for Image Denoising" (TIP 2017) - DnCNN
3. **Zhang et al.** - "Plug-and-Play Image Restoration with Deep Denoiser Prior" (TPAMI 2021) - DRUNet
4. **Zhu et al.** - "Denoising Diffusion Models for Plug-and-Play Image Restoration" (CVPR 2023) - DiffPIR

### Dataset
5. **Li et al.** - "Benchmarking Single-Image Dehazing and Beyond" (TIP 2019) - RESIDE Dataset

---

## Contact

**Author**: Zherui Zhang  
**Email**: zherui.zhang@mail.utoronto.ca  
**Institution**: University of Toronto
