CKDN — Image Quality Assessment with a Degraded Reference

PyTorch implementation of a Conditional Knowledge Distillation Network (CKDN) for blind image-quality assessment (IQA) of restored images when only a degraded reference is available.

Task	No- / degraded-reference image quality assessment
Paper	Learning Conditional Knowledge Distillation for Degraded Reference Image Quality Assessment
Framework	PyTorch (torch + patched, vendored timm)
License	MIT

---

Overview

In many real-world pipelines (restoration, super-resolution, compression), the pristine ground truth is unavailable — only a degraded reference exists. CKDN learns a model that predicts perceptual quality of a restored image relative to such a degraded reference, using a conditional knowledge-distillation strategy. The network is a two-branch ResNet-50 backbone:

• QSE (quality-sensitive encoder) — encodes the restored input,
• DTE (deep texture encoder) — encodes the degraded reference,

followed by shared comparison heads (csp/aux_csp) and per-head quality regressors producing scores in [0, 1]. The repository also exposes the intermediate feature extractors, which have been used for feature-space and MLP probes in the accompanying notebooks.

Repository layout

.
├── ckdn.py               # CKDN model definition (QSE/DTE two-branch ResNet-50)
├── train.py              # distributed training entry point
├── val.py                # SRCC/PLCC evaluation on the val split
├── predict_one_image.py  # inference: CLI + reusable IQA_CKDN class
├── train.sh / val.sh     # convenience launcher scripts
├── timm/                 # vendored + patched timm (IQA dataset & transforms)
├── images1/              # sample restored/degraded pairs for a quick demo
├── requirements.txt
└── README.md

Requirements

• Python 3.8+
• PyTorch 1.9+ (CUDA recommended)
• A NVIDIA GPU for training

timm is vendored under ./timm and locally patched for IQA datasets and paired transforms — do not pip install timm, or run the scripts from any directory other than the repository root.

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Quick start — inference

Predict the quality of a restored image against its degraded counterpart (sample pairs are provided under images1/restored and images1/degraded):

# restored/degraded live in sibling folders of the same name
python predict_one_image.py \
  --image images1/restored/3.png \
  --checkpoint /path/to/model_best.pth.tar

If the degraded image lives elsewhere, pass it explicitly:

python predict_one_image.py \
  --image /data/restored/3.png \
  --degraded /data/degraded/3.png \
  --checkpoint /path/to/model_best.pth.tar

The predicted score is printed to stdout (higher is better).

Programmatic API

from predict_one_image import IQA_CKDN

scorer = IQA_CKDN(checkpoint="model_best.pth.tar", device="cuda")
score = scorer.predict("images1/restored/3.png", "images1/degraded/3.png")

# Access intermediate representations for downstream probes:
dte = scorer.get_DTE_features("images1/degraded/3.png")
qse = scorer.get_QSE_features("images1/restored/3.png")
feat = scorer.extract_last_features("images1/restored/3.png", "images1/degraded/3.png")

The class falls back to CPU automatically when no GPU is available.

Training

The train script expects a dataset index (see Data format) and is launched via the torch.distributed.launch helper:

./train.sh ./data                # single GPU
NPROC_PER_NODE=8 ./train.sh ./data   # 8 GPUs

Checkpoints, logs, and args.yaml are written under ./output/<timestamp>-resnet-288/ (the --output flag overrides the base directory). model_best.pth.tar is saved whenever a new best validation metric is reached.

All hyper-parameters are configurable through the same flags used by the timm training scripts, e.g. --epochs, --lr, --batch-size, --sched, --img-size, --seed.

Validation

./val.sh /path/to/data /path/to/model_best.pth.tar

val.py reports SRCC and PLCC computed per content group and averaged across the validation set.

Data format

train.py/val.py read the data through the patched timm.data.Dataset (timm/data/dataset.py), which is tailored to a KADID-style degraded-reference dataset. Concretely, it requires:

1. train.txt and val.txt index files — one image_path,mos per line — living inside the dataset directory passed as DIR (the loader resolves DIR[1:] + "<phase>.txt" relative to the working directory),
2. image files named with the KADID convention *_<distortion>_<level>_*.bmp, where the distortion id/level tokens are used to sample pairing distortions at training time,
3. a ref_<name>.bmp (pristine) reference per distorted image.

Training is pair-wise: every sample is concatenated as [restored, ref, distorted, ref] → 12-channel input, and the targets are the corresponding MOS pairs. To train on your own data, adapt timm/data/dataset.py (and the loss expectations in train.py:loss_fn) to your naming and score conventions.

The per-phase .txt index files and raw dataset images are not committed to this repository.

Evaluation datasets

This codebase has been tested against publicly available IQA datasets — CSIQ and TID2013 — for full-reference and degraded-reference evaluation. Download them from their official sources and prepare the index files described above.

Implementation notes

• ckdn.model() bootstraps both branches from an ImageNet-pretrained ResNet-50 downloaded by torch.hub on first use.
• Checkpoints saved by DataParallel/DDP carry a module. key prefix; the inference scripts strip it automatically.
• The vendored timm/data/loader_iqa.py, paired_transforms_tv04.py, and config.py differ from upstream timm and are required by train.py/val.py — keep them in sync with ckdn.py if you ever refactor.