BasicSR

BasicSR (Basic Super Restoration) is a PyTorch toolbox from the XPixel Group for training and running image and video restoration models — super-resolution, denoising, deblurring, JPEG artifact removal, and face restoration. It bundles reference implementations of well-known research architectures such as EDSR, RCAN, SRResNet/SRGAN, ESRGAN, EDVR, BasicVSR, SwinIR, ECBSR, StyleGAN2, and DFDNet behind a single registry-driven training and testing pipeline.

Rather than shipping one model, BasicSR provides the scaffolding a restoration research project needs: YAML-configured experiments, a plugin registry for datasets/architectures/models/losses/metrics, distributed training support, and logging integrations (TensorBoard, Weights & Biases). It underpins downstream projects like Real-ESRGAN and GFPGAN and is commonly installed as a pip package to reuse individual archs, losses, or metrics inside custom training code.

What You Get

• Reference PyTorch implementations of major restoration architectures: EDSR, RCAN, SRResNet/SRGAN, ESRGAN, RRDBNet, EDVR, BasicVSR/BasicVSR++, SwinIR, ECBSR, StyleGAN2, and DFDNet.
• A decorator-based registry system (ARCH_REGISTRY, MODEL_REGISTRY, DATASET_REGISTRY, LOSS_REGISTRY, METRIC_REGISTRY) for plugging in custom components without touching the core training loop.
• train.py and test.py entry points that parse YAML experiment configs and drive the full training/validation/checkpointing lifecycle, including distributed (multi-GPU) training.
• Ready-to-run inference scripts for BasicVSR, BasicVSR++, DFDNet, ESRGAN, RIDNet, StyleGAN2, and SwinIR under inference/.
• Dataset utilities (paired/single image, REDS, Vimeo90K, FFHQ, LMDB-backed loaders) and degradation pipelines for building super-resolution training data.
• Custom CUDA extensions (deformable convolution, upfirdn2d, fused activation) compiled at install time for EDVR- and StyleGAN2-style architectures.

Common Use Cases

• Training a super-resolution or denoising model on a custom dataset using an existing architecture and a YAML config.
• Reusing BasicSR’s loss functions, metrics (PSNR/SSIM/NIQE/FID), or architecture definitions as building blocks inside a separate PyTorch project.
• Running inference with a pretrained restoration model (e.g. SwinIR or BasicVSR) via the provided inference scripts.
• Prototyping a new restoration architecture by registering it with BasicSR’s registry and reusing its training/validation loop.

Under The Hood

Architecture BasicSR is organized around five parallel registries (ARCH_REGISTRY, MODEL_REGISTRY, DATASET_REGISTRY, LOSS_REGISTRY, METRIC_REGISTRY) defined in basicsr/utils/registry.py; each package (archs/, models/, data/, losses/) auto-scans its own directory for files matching a naming convention (*_arch.py, *_model.py, etc.) and imports them so their @REGISTRY.register() decorators run at import time. train.py parses a YAML experiment config via basicsr/utils/options.py, builds the dataset and model through build_dataset/build_model (which just look up the registered class by the type string in the config), and drives a loop that owns iteration counting, EMA, LR scheduling, checkpointing, and distributed-training coordination. Models subclass BaseModel (basicsr/models/base_model.py), which defines the training/validation contract (feed_data, optimize_parameters, validation) that concrete models like SRModel and VideoBaseModel implement. Because nearly every arch, model, dataset, and loss in the codebase participates in this registry-and-inheritance scheme, changing BaseModel or the Registry class itself would ripple across most of the repository.

Tech Stack The project targets Python with PyTorch >=1.7 and torchvision as its core dependencies, plus opencv-python, Pillow, scikit-image, and scipy for image processing, lmdb for high-throughput dataset storage, PyYAML for config parsing, addict for attribute-style config dicts, and tb-nightly/optional Weights & Biases for training logs. Performance-critical ops (deformable convolution for EDVR, upfirdn2d and fused leaky-ReLU for StyleGAN2) are implemented as custom CUDA/C++ extensions compiled through setup.py using PyTorch’s CUDAExtension/CppExtension. Packaging is plain setuptools (no pyproject.toml), and releases are published to PyPI automatically via a GitHub Actions workflow triggered on version tags.

Code Quality The repository ships a tests/ suite (pytest) exercising individual architectures, models, datasets, losses, and metrics with real forward passes against small YAML-defined configs, and CI runs this alongside flake8, isort, yapf, and codespell on every push and pull request, backed by a .pre-commit-config.yaml. Error handling is minimal and mostly relies on assertions inside the registry and config parsing rather than explicit typed exceptions; there are no type hints or static type checking (no mypy, no py.typed marker). Naming is consistent throughout, following the <name>_arch.py / <name>_model.py suffix convention the registry scanning depends on.

API Design Using BasicSR as a library means writing a full YAML experiment config (network, dataset, training, and validation sections) rather than a few lines of Python, which gives real onboarding friction versus a typical import-and-call library — though the maintainers publish a companion BasicSR-Examples repo specifically to lower that barrier for custom projects. Individual pieces (an arch class, a loss function, a metric) can be imported and used directly without the training harness, and the extensive docs/ tree (Config, DatasetPreparation, HOWTOs, ModelZoo, TrainTest, FAQ, with Chinese translations) documents both the config-driven and library usage paths in detail.