A lightweight image-restoration pipeline for real-time low-light computer vision, using a GAN-trained U-Net teacher and knowledge distillation to compress the learned enhancement behavior into a small student model suitable for resource-constrained systems.
Tech: PyTorch · Computer Vision · U-Net · GANs · Knowledge Distillation · ResNet
Left to right: Ground truth | Model input (low ISO) | Model output (reconstruction)
Left to right: Ground truth | Model input (high ISO) | Model output (reconstruction)
| Model | Parameters | Size | Avg. Inference Time |
|---|---|---|---|
| Teacher U-Net | 38.1 M | 152.4 MB | 7.5 ms |
| Student U-Net | 2.4 M | 9.6 MB | 5.5 ms |
- Reduced the deployed model from 38.1 M to 2.4 M parameters through teacher-student knowledge distillation — a 15.9× reduction.
- Reduced the model footprint from 152.4 MB to 9.6 MB while retaining the teacher's overall reconstruction behavior.
- Achieved 5.5 ms average inference time per image in the project benchmark, demonstrating real-time feasibility for downstream vision pipelines.
- Built a preprocessing pipeline for the See-in-the-Dark (SID) dataset, converting RAW
.ARWcamera files into model-ready images and incorporating ISO and exposure time as additional input channels. - Evaluated enhancement quality using L1 reconstruction loss, SSIM, and PSNR; the teacher achieved 18.40 dB test PSNR.
Low-light images can significantly degrade downstream perception by hiding edges, textures, colors, and other features that computer-vision systems depend on.
The goal of this project was not only to reconstruct extremely dark images, but to make the enhancement stage small and fast enough to plausibly sit in front of a real-time perception system on platforms such as drones, mobile robots, vehicles, or embedded cameras.
The resulting pipeline separates training-time model capacity from deployment-time cost:
TRAINING
│
Dark Image ──> Teacher U-Net ──> Enhanced Image
▲
│ adversarial feedback
ResBlock Discriminator
│
▼
Realistic Reconstruction
DISTILLATION
│
Dark Image ─────────┼──────────────┐
▼ ▼
Teacher U-Net Student U-Net
│ │
└──── L1 ──────┘
DEPLOYMENT
│
Dark Image ──> 2.4M-parameter Student U-Net ──> Enhanced Image
The teacher is an encoder-decoder U-Net designed to retain both global scene structure and fine image detail.
Key design choices:
- Residual encoder blocks with strided convolutions for downsampling
- Bilinear upsampling + convolution in the decoder
- Learned attention-gated skip connections between corresponding encoder and decoder stages
- 5-channel input: RGB plus image-wide ISO and exposure-time channels
- Training resolution of 256 × 256
The skip connections preserve local image structure while the bottleneck captures broader lighting and scene context.
Pixel-level reconstruction alone can reward blurry or visually implausible outputs. To improve perceptual realism, the teacher was trained with a separate ResBlock CNN discriminator using hinge adversarial loss.
The generator objective combines:
Generator Loss = α · L1 + β · SSIM Loss + γ · Adversarial Loss
The final experiment used reconstruction, structural-similarity, and adversarial terms together, with generator warmup and discriminator cooldown logic to stabilize GAN training.
By the end of training, the discriminator and generator faced similar loss values, indicating convergence of the two to a stable state.
The full teacher is useful during training but unnecessarily large for deployment. A smaller U-Net student therefore learns directly from the frozen teacher's enhanced outputs.
Both networks receive the same dark image:
Teacher output = T(x)
Student output = S(x)
Distillation loss = L1(T(x), S(x))
The student keeps the main encoder-decoder design but uses substantially narrower layers and trains at 128 × 128 resolution.
This reduced the deployed network to 2.4 M parameters / 9.6 MB, compared with 38.1 M parameters / 152.4 MB for the teacher.
The model was evaluated on paired short- and long-exposure images from SID. The reconstruction pipeline recovered scene structure and visible foreground information from heavily underexposed inputs, although performance degraded for the darkest exposures and high-noise images.
Examples can be seen on the top of this README.
Left to right: Ground truth | Teacher reconstruction | Student reconstruction
Despite using roughly 1/16 the parameters, the student reproduced the teacher's overall colors, scene structure, and reconstruction patterns closely on held-out images.
This was the primary deployment result of the project: knowledge distillation substantially reduced memory cost while preserving the behavior learned by the larger GAN-trained network.
The project uses the See-in-the-Dark (SID) Sony dataset, containing:
- 5,094 short-exposure RAW images
- 424 long-exposure reference images
- Indoor scenes captured at approximately 0.03–0.3 lux
- Outdoor scenes captured at approximately 0.2–5 lux
A custom preprocessing pipeline converts Sony .ARW files into PNG images before training.
Each model input contains:
[R, G, B, ISO, Exposure Time]
ISO and exposure time are expanded into two additional image-sized channels so that the network can condition reconstruction on camera settings.
.
├── models/ # Teacher, student, discriminator, and model components
├── utils/ # Metrics and supporting utilities
├── preprocess.py # RAW-to-PNG SID preprocessing
├── preprocess.yaml # Dataset preprocessing paths
├── config.yaml # Training configuration
├── run.py # Training entry point
├── requirements.txt
└── README.md
Adjust the structure above if the repository directory names differ.
Create and activate a Conda environment:
conda create --name lowlight-unet python
conda activate lowlight-unet
pip install -r requirements.txtA CUDA-capable GPU is recommended for reproducing training. The original teacher experiment took approximately 6 hours for 50 epochs on an NVIDIA RTX 3070.
Download the Sony subset of the SID dataset and configure the raw and processed dataset locations in:
preprocess.yaml
Convert the RAW files with:
python preprocess.pyThe expected dataset structure is:
SIDDataset/
├── Sony/
│ ├── long/
│ └── short/
├── Sony_PNG/
│ ├── long/
│ └── short/
├── Sony_test_list.txt
├── Sony_train_list.txt
└── Sony_val_list.txt
Train the teacher model:
python run.pyUse a custom configuration:
python run.py --config config.yamlResume from a PyTorch Lightning checkpoint:
python run.py --resume <checkpoint.ckpt>Train the knowledge-distilled student:
python run.py --studentPyTorch Lightning saves experiment logs and reconstruction samples under logs/.
The project also exposed several important tradeoffs in low-light restoration:
- Reconstruction quality degraded significantly at the shortest exposure times, where the input contained very little recoverable visual information.
- High-ISO images remained difficult to denoise, partly because high-ISO examples represented a smaller fraction of the training data.
- SSIM was less effective as a direct optimization term than expected and may be better used primarily as an evaluation metric.
- Adversarial training improved the modeling objective but introduced optimization instability and occasional color artifacts.
- The strongest result was the compression pipeline: the student retained the teacher's reconstruction characteristics with a fraction of the model size.
These findings motivate future work with perceptual losses, improved adversarial objectives, better balancing of exposure conditions, and deployment on embedded GPU hardware.
The project was designed around environments where lighting cannot be assumed to be ideal. A small enhancement network could act as a preprocessing stage before:
- visual odometry / SLAM
- object detection
- obstacle recognition
- autonomous navigation
- other camera-based robotic perception tasks
The 9.6 MB student model and millisecond-scale inference time make this architecture substantially more practical for latency-sensitive vision pipelines than the full training-time model.