nvRestorer

nvRestorer is an experimental, GPU‑centric video mosaic restorer built around:

• PyNvVideoCodec (NVDEC / NVENC) for GPU decode / encode
• A YOLO‑seg mosaic detector (Lada detection models)
• A BasicVSR++‑style mosaic restorer (ported from Lada’s mmagic pipeline)
• A simple Python CLI so you can process individual video files end‑to‑end

The goal is to mirror Lada’s restoration behaviour as closely as possible, but inside a fully GPU‑accelerated pipeline that keeps decode, detect, restore, and encode on the GPU as much as practical.

⚠️ This is work‑in‑progress. The CLI, options and internal APIs may change.

---

Status – 0.0.5

Implemented:

• NVDEC decode via PyNvVideoCodec (video.Decoder)
• GPU RGB frame pipeline
• YOLO mosaic detection with masks (detector.Detector)
• Temporal buffer for per‑ROI context (restorer.TemporalBuffer)
• BasicVSR++ restorer using Lada weights (restorer.RealRestorer)
• ROI‑based paste‑back with soft, mask‑driven blending
• NVENC encode via PyNvVideoCodec (video.Encoder)
• Debug PNG export for the first N frames (optional)

Known limitations:

• Some extreme cases (large merged ROIs, mosaics close to skin tone) still produce artifacts.
• Performance tuning is still basic (single process, single CUDA stream).
• Only single‑file input/output is supported right now (no batch dirs yet).

---

Repository Layout

At a high level:

nvRestorer/
├── cli/              # CLI entry point and pipeline wiring
│   └── main.py
├── detector/         # YOLO-based mosaic detector wrapper
│   ├── core.py
│   └── yolo.py
├── restorer/         # Restorer base + pseudo + real (BasicVSR++)
│   ├── base.py
│   ├── pseudo_restorer.py
│   └── real_restorer.py
├── video/            # NVDEC / NVENC wrappers (decode/encode)
│   ├── decoder.py
│   └── encoder.py
├── models/           # Model definitions
│   ├── lada/
│   │   ├── basicvsr_plusplus_net.py
│   │   └── ...
│   └── basicvsrpp.py
├── utils/            # Color, letterbox, misc helpers
├── __init__.py
└── pyproject.toml / requirements.txt / README.md

The main execution path is:

cli/main.py → pipeline.run_file()
            → video.Decoder / detector.Detector / restorer.RealRestorer / video.Encoder

---

Pipeline Overview

For each input video:

1. Decode (NVDEC)
   video.Decoder uses PyNvVideoCodec to decode frames on the GPU. Surfaces are converted to RGB tensors that stay on the GPU as much as possible.

2. Detect mosaics (YOLO‑seg)
   detector.Detector wraps a YOLO‑seg model (Lada’s mosaic detection model).
   For each frame it produces:

   • bounding boxes [x1, y1, x2, y2]
   • binary masks (mosaic pixels)

   restorer.BaseRestorer.build_rois_for_frame turns these into Roi objects, each containing:

   • the ROI box (frame coordinates)
   • an ROI‑sized binary mask

3. Temporal buffer
   restorer.TemporalBuffer keeps a sliding window of RGB frames on the GPU.
   For frame t, RealRestorer can pull frames [t‑R … t … t+R] where R = temporal_radius.

4. Restore ROIs (RealRestorer / BasicVSR++)

   For each ROI in each frame:

   • Crops the ROI patch from each frame in the temporal window.
   • Normalizes to [0, 1] or [-1, 1] (configurable).
   • Pads the patch symmetrically to at least 256×256 and multiples of 32 so that BasicVSR++’s internal 4× downsampling sees a ≥64×64 low‑res map.
   • Runs BasicVSRPlusPlusNet (Lada‑compatible weights) on the padded sequence.
   • Crops the central, unpadded region corresponding to the original ROI patch.
   • Builds a soft blend mask from the ROI’s binary mask (Torch port of Lada’s create_blend_mask).
   • Alpha‑blends the restored patch back into the original frame at [y1:y2, x1:x2].

5. Encode (NVENC)
   video.Encoder wraps PyNvVideoCodec’s NVENC path to encode the final frames back to a file, using a configurable codec and options.

---

Requirements

• OS:

  • Windows 10/11 tested
  • Linux should work if you can install PyTorch + CUDA + PyNvVideoCodec.

• Python:

  • 3.10+ recommended

• GPU / CUDA:

  • NVIDIA GPU with NVENC + NVDEC support (e.g. Pascal and newer)
  • CUDA‑enabled PyTorch (e.g. torch 2.x + cu12x)

• Libraries (rough sketch):

  • torch
  • PyNvVideoCodec
  • numpy
  • Pillow
  • ultralytics or equivalent YOLO library (depending on final detector wiring)
  • standard Python libs

See pyproject.toml or requirements.txt for the exact list used in this snapshot.

You will also need the Lada models:

• lada_mosaic_detection_model_v3.1_accurate.pt
• lada_mosaic_restoration_model_generic_v1.2.pth

You can place them anywhere; the CLI takes explicit paths.

---

Installation

Clone and set up a virtual environment:

git clone https://github.com/seatv/nvRestorer.git
cd nvRestorer

python -m venv .venv
# Windows:
.\.venv\Scripts\activate
# Linux/macOS:
# source .venv/bin/activate

pip install -r requirements.txt
# or, if packaged:
# pip install .

---

CLI Usage

The CLI entry is the nvRestorer.cli module. The exact options may evolve, but the current pattern is:

python -m nvRestorer.cli INPUT OUTPUT [options]

Example (Windows, PowerShell):

python -m nvRestorer.cli `
  D:\ToClean\2K-Test.mp4 `
  D:\Results\2K-Test-Real.mp4 `
  --restorer real `
  --restorer-model D:\Models\lada\lada_mosaic_restoration_model_generic_v1.2.pth `
  --detector-model D:\Models\lada\lada_mosaic_detection_model_v3.1_accurate.pt `
  --batch-size 60 `
  --restorer-max-size 512 `
  --restorer-normalize 0-1

Example (Linux):

python -m nvRestorer.cli \
  /data/in/2K-Test.mp4 \
  /data/out/2K-Test-Real.mp4 \
  --restorer real \
  --restorer-model /models/lada_mosaic_restoration_model_generic_v1.2.pth \
  --detector-model /models/lada_mosaic_detection_model_v3.1_accurate.pt \
  --batch-size 60 \
  --restorer-max-size 512 \
  --restorer-normalize 0-1

Key options (current snapshot)

Note: check cli/main.py for the authoritative list and defaults.

• INPUT / OUTPUT
  Positional arguments: input video path and output video path.

• --device-index N
  Select CUDA device index (default: 0).

• --batch-size N
  Number of frames per detection/restoration batch.

• --restorer {pseudo,real}

  • pseudo: cheap patch‑based restorer that doesn’t use BasicVSR++.
  • real: full BasicVSR++ restorer (Lada weights).

• --restorer-model PATH
  Path to the BasicVSR++ restoration checkpoint (*.pth).

• --detector-model PATH
  Path to the YOLO mosaic detection model.

• --restorer-max-size N
  Skip ROIs larger than N × N pixels (safety cutoff).

• --restorer-normalize {0-1,-1-1}
  Pixel normalization scheme for the restorer:

  • 0-1 → values in [0,1]
  • -1-1 → values in [-1,1] (Lada‑style)

• Debug (dev use)
  Depending on wiring in cli/main.py, you can pass:

  • --debug-frames N – save debug PNG tiles for first N frames.
  • --debug-dir PATH – where to write those PNGs.

The debug PNGs, when enabled, include (per ROI):

• *_1_original.png – original ROI patch
• *_2_restored.png – BasicVSR++ output
• *_3_mask.png – ROI mask
• *_4_final.png – final blended patch

---

Roadmap

Short‑term ideas:

• ROI hygiene

  • Tighten ROIs to mask‑derived bounds instead of raw detector boxes.
  • Optionally split multi‑blob masks into per‑blob ROIs to avoid very wide stripes.

• Performance

  • Overlap decode / detect / restore / encode using multiple threads or processes.
  • Tune batch sizes and temporal radius for different GPU classes.

• CLI & UX

  • Add directory mode (batch processing over a folder).
  • Improve validation and error messages for missing models / bad codecs.

Medium‑term:

• More robust colour / gamma handling to better match Lada output.
• Support CPU / XPU backends as a fallback where CUDA isn’t available.
• Better metrics / profiling hooks for performance tuning.

---

Acknowledgements

This project draws heavily from:

• Lada – for the detection and restoration models and the original mmagic‑based pipeline.
• BasicVSR++ – for the underlying video restoration architecture.

Please check the upstream projects for full training code, original implementations, and model weights.

---

Disclaimer

This software is provided for research and educational purposes only.

You are responsible for ensuring that your use of the software and models complies with all applicable laws and respects the rights and consent of all people depicted in any media you process.

License

nvRestorer is distributed under the terms of the GNU Affero General Public License v3.0 (AGPL-3.0).

This project adapts ideas and logic from the Lada mosaic restoration project, which is also licensed under AGPL-3.0.