dronefreak/yolo26-rgb-s

yolo26-rgb: RGB Image Restoration (YOLO26s)

YOLO26's dense depth-estimation head, repurposed to output 3-channel RGB instead of 1-channel depth. YOLO26 already ships a depth-estimation head, itself a dense, full-resolution regression task; this model swaps that head for an RGB one and asks whether the decoder's inductive bias, tuned for real-time detection speed and then adapted once for depth, transfers to single-image deraining any better than a classification-pretrained backbone repurposed as a UNet encoder (ClearView's ResNet-UNet family). Trained and evaluated on ClearView's mixed-domain recipe and 10-test-set protocol, for numbers directly comparable to ClearView's own model zoo, not a separate benchmark. See Model Comparison below for how this scale specifically does against that baseline.

Best PSNR/SSIM of the two released scales (30.95 avg, 9 rain-only sets), still fast (92.2 qps). The default choice unless a smaller footprint than its 12.13M params is required, in which case use n.

---

Architecture

Backbone and PAN-FPN neck are unmodified YOLO26. RGBHead replaces the 1-channel depth head with a 3-channel restoration head: full input-resolution output, two skip connections from the stride-2/4 backbone layers into the decoder tail, a global residual (predict a correction, add it to the input), and LayerNorm in the head's own conv blocks while the backbone/neck stay on BatchNorm (folds into conv at TensorRT export; keeps Ultralytics' full YOLO26 pretrained zoo loadable). Full write-up and the RGBHead-internals figure: github.com/dronefreak/yolo26-rgb#architecture.

---

Quickstart

pip install git+https://github.com/dronefreak/yolo26-rgb.git

import torch
from yolo26_rgb import YOLO26RGB

model = YOLO26RGB.from_pretrained("dronefreak/yolo26-rgb-s")  # downloads yolo26s-rgb.pt, returns eval mode

# image_tensor: (1, 3, H, W) float in [0, 1]. The output is a residual added to the
# input (NAFNet/Restormer-style), not bounded to [0, 1] - clamp it before saving.
output = model(image_tensor).clamp(0, 1)

Full image-to-PNG example (load an image, run, write the result): see the main repo README.

---

Training Data

5 sources combined via ClearView's --mix-config (recipe), oversampling the real-world sources 2x:

Source	Type	Weight	Pairs
Rain13K [1]	Synthetic	1.0	13,711
DDN-Data / Rain1400 [4]	Synthetic	1.0	12,600
SPA-Data [5]	Real-world	2.0	6,385
RealRain-1k-H [6]	Real-world	2.0	784
RealRain-1k-L [6]	Real-world	2.0	784

~62% synthetic / ~38% real by effective sampling weight (before oversampling: ~77%/23% by raw pair count).

Checkpoint selection uses a separate blended validation set (recipe): SPA-Data val (capped to 150 of 1,000 pairs so it can't dominate), RealRain-1k-H/L validation (112 each), and Rain100L (100) as a synthetic sanity anchor.

---

Detailed Test-Set Metrics

Full per-dataset breakdown across all 6 tracked metrics, computed on each source's own held-out test/eval split (not the blended validation set used for checkpoint selection during training).

Test Set	Domain	PSNR	SSIM	MAE	MSE	Rain Removal Rate	NIQE
SPA-Data_val [5]	Real-world	38.66	0.966	0.0065	0.00036	0.323	8.93
Rain100L [2]	Synthetic	30.41	0.893	0.0185	0.00118	0.171	19.60
Rain100H [2]	Synthetic	23.10	0.566	0.0540	0.00553	0.423	54.59
Test100 [3]	Synthetic	26.03	0.775	0.0443	0.00397	0.322	17.31
Test1200 [3]	Synthetic	29.01	0.814	0.0314	0.00220	0.293	13.21
Test2800 [4]	Synthetic	29.97	0.861	0.0240	0.00112	0.290	761.51
DDN-Data [4]	Synthetic	30.24	0.866	0.0234	0.00107	0.280	1223.47
RealRain-1k-H [6]	Real-world	34.19	0.911	0.0139	0.00074	0.535	9.37
RealRain-1k-L [6]	Real-world	36.92	0.950	0.0098	0.00041	0.498	7.77
AllWeather (rain+fog) [7]	Cross-domain (stress)	13.54	0.566	0.1918	0.05658	0.067	273.51
Average (10 sets)		29.21	0.8167
Average (9 rain-only)		30.95

Metric definitions:

• PSNR (Peak Signal-to-Noise Ratio, dB): pixel-level fidelity between the restored output and ground truth. Higher is better.
• SSIM (Structural Similarity Index, 0-1): perceptual similarity based on luminance, contrast, and structure. Higher is better.
• MAE (Mean Absolute Error, normalized [0,1] pixel space): average per-pixel absolute intensity difference. Lower is better.
• MSE (Mean Squared Error, normalized [0,1] pixel space): average per-pixel squared intensity difference, the term PSNR is derived from. Lower is better.
• Rain Removal Rate: ClearView-specific metric. Compares the Sobel-gradient ("high-frequency") residual energy of the output vs. ground truth against that of the input vs. ground truth: 1 - (residual after) / (residual before). 1.0 means perfect rain removal, 0.0 means no change, negative means the model added more high-frequency error than it removed. Higher is better.
• NIQE (Natural Image Quality Evaluator): no-reference perceptual quality score, refit per test set from that set's own clean images, so NIQE values are only comparable within the same row, not across rows. Lower is better within a given row.

AllWeather (rain+fog) is an explicit out-of-domain stress test, fog was not in the training mix and this model is not expected to handle it well; excluded from the 9-rain-only average above, same convention ClearView itself uses to rank its own model zoo.

---

Model Comparison

How this model compares to ClearView's model zoo and the other released yolo26-rgb scale (this row in bold), ranked by average PSNR across the 9 rain-only test sets, same test sets and evaluation protocol for every model. The three deployment columns are TensorRT 11.2.1, fp16, 1920x1080 (padded to 1088x1920), batch size 1, on an RTX 4070 SUPER (12GB), CUDA 12.8, cuDNN 9.8.0, Ubuntu 24.04.4 LTS; ClearView baseline figures are its own published benchmark on the identical GPU/TensorRT version.

Rank	Model	Params	Avg PSNR (9 rain-only)	Engine (fp16)	Latency (fp16)	Throughput (fp16)
1	Restormer [8]	15.3M	35.10	OOM*	OOM*	OOM*
2	NAFNet (Large) [10]	116M	34.16	3.86 GB	1938.4 ms	0.52 qps
3	NAFNet (Mid) [10]	14.3M	33.97	1.85 GB	305.4 ms	3.3 qps
4	Restormer-Small [8]	2.3M	31.98	7.2 GB	211.2 ms	4.7 qps
5	UNet (Vanilla) [9]	21.5M	31.74	44 MB	36.6 ms	27.3 qps
6	NAFNet (Small) [10]	1.1M	31.15	911 MB	37.2 ms	26.9 qps
-	yolo26_rgb_s	12.13M	30.95	1.53 GB	10.85 ms	92.2 qps
-	yolo26_rgb_n	5.25M	30.83	1.27 GB	9.20 ms	108.6 qps
7	ResNet50-UNet [11]	73.3M	30.63	315 MB	30.2 ms	33.1 qps
8	ResNet34-UNet [11]	24.5M	30.45	217 MB	10.5 ms	94.9 qps
9	ResNet18-UNet [11]	14.4M	30.23	197 MB	9.07 ms	110.3 qps

* Restormer's full-size checkpoint doesn't survive TensorRT conversion at this resolution: ONNX export itself OOMs at fp32 (over 12GB during tracing), and even at fp16 the TensorRT build fails, needing roughly 14.4GB of activation/scratch memory to fuse its attention path on a 12GB card.

---

Use Cases

Good for: real-time or edge-constrained deraining where a compact, TensorRT-exportable backbone matters; research into whether depth-pretrained detection/segmentation backbones transfer to restoration tasks better than classification-pretrained ones.

Limitations: deraining is partial, not complete, residual streaks survive up close, worse on dense rain over flat/low-texture backgrounds; AllWeather (rain+fog) is a stark outlier (~13.5 dB vs. 22+ dB on every other test set), an explicit out-of-domain stress test, not representative of typical performance; trails Restormer/NAFNet by a wide margin on PSNR, those are accuracy-maximizing architectures with no real-time constraint, never the target this project was built to beat, see ClearView's ResNet-UNet family for the comparison that is.

---

License: AGPL-3.0

AGPL-3.0, inherited from depending on AGPL-licensed Ultralytics YOLO26 source (Ultralytics requires the entire derivative work, including training your own model from scratch, to stay AGPL-3.0 or use their Enterprise license, see their licensing page). This includes AGPL's network-use clause (section 13): if you run this model as part of a network-accessible service, you're obligated to offer that service's complete corresponding source to its users, not just link back here. Not affiliated with or endorsed by Ultralytics, this repurposes their published architecture for a different task, no claim of ownership over the original YOLO26 design.

---

Citation

@software{saksena2026yolo26rgb,
  author = {Saksena, Saumya Kumaar},
  title = {yolo26-rgb: YOLO26's depth head repurposed for RGB image restoration},
  year = {2026},
  url = {https://github.com/dronefreak/yolo26-rgb}
}

Architecture and datasets this model builds on:

References

1. Fu et al. and others, Rain13K composite (Rain800/Rain100H/L/Rain14000/Rain12), standard MPRNet/Restormer training set.
2. Yang et al., Deep Joint Rain Detection and Removal from a Single Image, CVPR 2017 (Rain100H/L).
3. Zhang & Patel / Zhang, Sindagi & Patel (Test100 / Test1200).
4. Fu et al., Removing Rain from Single Images via a Deep Detail Network, CVPR 2017 (DDN-Data / Rain1400 / Test2800).
5. Wang et al., Spatial Attentive Single-Image Deraining with a High Quality Real Rain Dataset, CVPR 2019 (SPA-Data).
6. Li et al., RealRain-1k, arXiv:2206.05514, 2022.
7. Li et al., Heavy Rain Image Restoration, CVPR 2019 (AllWeather rain+fog / Outdoor-Rain).
8. Zamir et al., Restormer: Efficient Transformer for High-Resolution Image Restoration, CVPR 2022, arXiv:2111.09881.
9. Ronneberger, Fischer & Brox, U-Net: Convolutional Networks for Biomedical Image Segmentation, MICCAI 2015, arXiv:1505.04597.
10. Chen, Chu, Zhang & Sun, Simple Baselines for Image Restoration, ECCV 2022, arXiv:2204.04676 (NAFNet).
11. He, Zhang, Ren & Sun, Deep Residual Learning for Image Recognition, CVPR 2016, arXiv:1512.03385 (ResNet, the encoder backbone for the ResNet18/34/50-UNet baselines).
12. Jocher et al., Ultralytics YOLO26: Unified Real-Time End-to-End Vision Models, 2026, arXiv:2606.03748 (this model's backbone/neck, and the depth-estimation head it's derived from).

Full bibtex for the YOLO26 architecture: main repo README. ClearView's own citation: its README.

---

Model Card Authors / Contact

Saumya Kumaar Saksena, GitHub Issues