YOLO26-RGB: From Depth to Drizzle — Repurposing YOLO26’s Backbone for Image Deraining

# Repurposing a Depth-Estimation Backbone for Image Deraining: Architecture, Experiments, and Results

## Introduction

Image deraining — the task of removing rain streaks and artifacts from wet or rainy photographs — is a dense-pixel regression problem that shares architectural DNA with depth estimation. Both require the model to produce a full-resolution, per-pixel output conditioned on local and global spatial features. This overlap raised an interesting question: could a backbone and feature-pyramid neck trained for depth estimation serve as a strong initialization for a deraining model, compared with training the same architecture from scratch?

The project explores this by taking the pretrained weights from a depth-estimation model and transplanting them into a freshly designed restoration decoder, then benchmarking the result against both random initialization and established deraining methods.

## Architecture Overview

### What Was Carried Over

The foundation comes from a modern convolutional backbone paired with a feature pyramid network. Specifically:

– **Backbone:** A CSPDarknet-style architecture that extracts multi-scale features through hierarchical downsampling.
– **Neck:** A PAN-FPN pathway that fuses features across scales, producing a multi-resolution feature pyramid (typically referred to as P3, P4, and P5).
– **Feature Fusion:** The multi-scale decoder projects all pyramid levels to a common channel width, then progressively upsamples and adds them in a top-down fashion (P5 → P4 → P3). This fusion mechanism is general-purpose and not specific to depth estimation.

These components transfer cleanly because depth regression and image restoration are structurally similar tasks — both demand dense, pixel-level predictions across multiple spatial scales.

### What Was Modified or Added

Several deliberate changes were made to adapt the transferred features for deraining:

1. **New Restoration Decoder:** A dedicated decoding head was written from scratch, replacing the original single-channel depth output. This was not simply a re-pointed depth head; it was an entirely new module designed for restoration.

2. **Full-Resolution Reconstruction Tail:** Because deraining requires pixel-exact output (unlike depth estimation, which can operate at reduced resolution), the tail of the network continues processing all the way back to the input resolution.

3. **Skip Connections:** Direct connections from the stride-2 and stride-4 layers of the backbone feed into the reconstruction tail. This ensures that fine spatial detail — edges, textures, small rain streaks — has a path that bypasses the heavily downsampled bottleneck.

4. **Residual Prediction:** Rather than predicting the clean image directly, the head outputs a correction signal that is added to the rainy input. This residual formulation, common in modern restoration networks, makes the learning target easier and stabilizes training.

5. **Normalization Choices:** The decoder blocks within the new head use Layer Normalization, while the backbone and neck retain Batch Normalization. This decision keeps the model loadable from the broader pretrained model zoo (which uses Batch Norm), and it also allows the Batch Norm layers to fold into adjacent convolution operations during TensorRT export.

## Training Setup and Evaluation Protocol

The model was trained using an external deraining library that provides a mixed synthetic-and-real rain dataset, a Charbonnier loss function, and a standardized 10-test-set evaluation protocol. Two model scales were released:

| Scale | Parameters |
|——-|———–|
| Nano | 5.25 M |
| Small | 12.13 M |

The critical controlled experiment compared two initialization strategies at the nano scale: loading the depth-pretrained backbone and neck weights versus starting from random initialization. Both conditions used the same architecture, the same training recipe, and a fixed budget of 100 epochs.

## The Transfer Experiment: Depth vs. Random Initialization

The headline result from the controlled initialization experiment is striking:

| Initialization | Avg PSNR (10 sets) | Avg SSIM | Test Sets Won |
|—————|——————-|———-|—————|
| Random | 27.45 | 0.807 | 0 / 10 |
| Depth-Pretrained | 27.94 | 0.813 | 10 / 10 |
| **Delta** | **+0.48 dB** | **+0.006** | — |

The depth-initialized model won every single test set. The delta appeared early — by 20 epochs the gap was already roughly +0.49 dB — and did not close with extended training. A single-epoch check showed the difference was not a convergence-speed artifact; the advantage was present from the start and persisted throughout.

These numbers should be interpreted with appropriate context: they are averages over 10 test sets including AllWeather (rain and fog combined), and they come from a 100-epoch run. The released models were trained for longer and report higher PSNR values on rain-only subsets.

An important caveat: this experiment does not explain *why* depth pretraining helps. It could be that depth supervision teaches the backbone useful geometric and spatial-structure representations that transfer to restoration, or it could simply be that the depth checkpoint is a particularly strong pretrained representation. The evidence here is that the representation is a better starting point, full stop.

## Accuracy Benchmarks

When evaluated on ClearView’s standard ranking convention (average PSNR over nine rain-only test sets), the released models compare favorably with a range of established deraining architectures:

| Model | Parameters | Avg PSNR |
|——-|———–|———-|
| Restormer | 15.3 M | 35.10 |
| NAFNet-Large | 116 M | 34.16 |
| NAFNet-Mid | 14.3 M | 33.97 |
| Restormer-Small | 2.3 M | 31.98 |
| UNet | 21.5 M | 31.74 |
| NAFNet-Small | 1.1 M | 31.15 |
| **yolo26_rgb_s** | **12.13 M** | **30.95** |
| **yolo26_rgb_n** | **5.25 M** | **30.83** |
| ResNet50-UNet | 73.3 M | 30.63 |
| ResNet34-UNet | 24.5 M | 30.45 |
| ResNet18-UNet | 14.4 M | 30.23 |

The YOLO-derived models punch well above their weight in terms of parameter count. NAFNet-Small achieves higher PSNR at 1.1 million parameters, but it runs roughly four times slower (26.9 queries per second), meaning it is not Pareto-optimal on every axis. The story here is specifically about the real-time, YOLO-derived operating point.

## Deployment and Real-World Performance

All benchmarks below use TensorRT fp16 inference at 1920×1080 resolution with batch size one on an RTX 4070 SUPER with 12 GB of VRAM.

The pairwise comparisons against the ResNet-UNet family highlight the efficiency advantage:

– **yolo26_rgb_s** (12.13 M, 30.95 dB) runs at 92.2 queries per second, matching the speed of **ResNet34-UNet** (24.5 M, 30.45 dB, 94.9 qps) while using roughly half the parameters and delivering +0.5 dB higher PSNR.
– **yolo26_rgb_n** (5.25 M, 30.83 dB) achieves 108.6 qps at the same throughput as **ResNet18-UNet** (14.4 M, 30.23 dB, 110.3 qps), with about one-third the parameters and +0.6 dB better quality.
– Both YOLO-derived models are approximately three times faster than **ResNet50-UNet** (73.3 M, 30.63 dB, 33.1 qps), while also scoring higher on PSNR.

One notable finding: the top-ranked model by PSNR (Restormer) could not be exported to TensorRT at 1080p on this GPU. The attention pathway required approximately 14.4 GB of temporary workspace, exceeding the 12 GB available. This illustrates a practical tradeoff between architectural sophistication and deployment feasibility.

## What This Shows — And What It Does Not

**Demonstrated:**
– Depth-pretrained backbone and neck weights provide a better initialization for deraining than random initialization, winning all 10 test sets in a controlled experiment with a +0.48 dB average PSNR advantage.
– The trained models occupy an attractive quality-latency Pareto frontier relative to ResNet-UNet baselines, offering higher PSNR at lower parameter counts and comparable or better throughput.

**Supported:**
– The CSPDarknet + PAN-FPN architecture, originally designed for depth estimation, transfers effectively to a dense restoration task with modest architectural modifications.

**Not Demonstrated:**
– Whether depth pretraining outperforms classification pretraining for restoration (no classification-pretrained baseline was tested).
– The mechanistic reason depth features are beneficial for deraining.
– Generalization beyond single-task deraining on rainy images.

## Limitations

Deraining remains a partially solvable problem. Faint streaks can survive close inspection, and dense rain over flat, low-texture backgrounds remains the worst-case scenario. The AllWeather benchmark (combining rain and fog) is out of distribution — both YOLO-derived models and all ClearView baselines land around 13.5 dB there, so it is excluded from the ranking. The project covers one task at two scales; it is not a general-purpose restoration model.

This was a solo side project released under the AGPL-3.0 license (inherited from the underlying YOLO framework). It is not affiliated with the organization that developed the base YOLO architecture.

—

## Frequently Asked Questions

**Q: Why use a depth-estimation backbone for deraining instead of a classification backbone?**
A: Depth estimation and image restoration are both dense, per-pixel regression tasks that require rich multi-scale spatial representations. Classification backbones are optimized for global feature extraction and are typically paired with upsampling decoders that lose fine spatial detail. A depth backbone, by contrast, is already trained to preserve pixel-level spatial structure throughout its hierarchy, which aligns more closely with the needs of deraining.

**Q: How much of the original depth model was actually reused?**
A: The backbone and neck — 468 tensors — were transferred with identical weights. Only the new restoration decoder head was randomly initialized and trained from scratch.

**Q: What is the significance of the +0.48 dB PSNR delta?**
A: In image quality metrics, even small PSNR differences can be perceptible, especially at the margins. A consistent +0.48 dB across all 10 test sets — with the gap appearing as early as 20 epochs and not closing over 100 epochs — indicates a genuine, stable advantage from the depth initialization rather than a transient or dataset-specific effect.

**Q: Can these models run on GPUs other than the RTX 4070 SUPER?**
A: The TensorRT fp16 benchmarks were reported on that specific GPU. The models themselves are architecture-agnostic and should export to TensorRT or other inference runtimes on compatible hardware. Performance numbers will vary based on GPU capability, VRAM, and batch size.

**Q: Does the nano model really outperform much larger models in speed?**
A: Yes. At 5.25 million parameters and 108.6 queries per second, the nano model delivers higher PSNR than the ResNet18-UNet (14.4 M, 30.23 dB, 110.3 qps) with about one-third the parameters and comparable throughput. Against the heavier ResNet50-UNet, the speed advantage is roughly threefold.

**Q: Why not simply use NAFNet-Small if it has better PSNR?**
A: NAFNet-Small achieves 31.15 dB versus 30.83 dB for the nano model — a real quality gap, but it runs at 26.9 queries per second, which is about four times slower. The choice depends on whether the priority is peak quality or real-time throughput. The YOLO-derived models are optimized for the latter.

**Q: Is this project affiliated with the developers of YOLO?**
A: No. This is an independent, solo project. The AGPL-3.0 license is inherited from the open-source framework on which the base architecture is built.

—

## Conclusion

This project demonstrates that a depth-estimation backbone can serve as an effective initialization for image deraining, yielding a measurable and consistent quality advantage over random initialization in a controlled experiment. The resulting models achieve a competitive balance of PSNR, parameter count, and inference speed, particularly when compared against the ResNet-UNet family. While the approach does not explain the underlying mechanism by which depth pretraining helps restoration, it does provide strong empirical evidence that the transferred representation is a superior starting point. The work highlights the value of cross-task pretraining in computer vision and opens the door to further exploration of which pretrained features are most useful for dense restoration tasks.

Thank you for reading