Showing preview only (464K chars total). Download the full file or copy to clipboard to get everything.
Repository: Zj-BinXia/KDSR
Branch: master
Commit: c7c9e8110654
Files: 135
Total size: 425.7 KB
Directory structure:
gitextract_r69x67a9/
├── KDSR-GAN/
│ ├── .gitignore
│ ├── .pre-commit-config.yaml
│ ├── CODE_OF_CONDUCT.md
│ ├── Metric/
│ │ ├── DISTS/
│ │ │ ├── DISTS_pytorch/
│ │ │ │ ├── DISTS_pt.py
│ │ │ │ └── weights.pt
│ │ │ ├── DISTS_tensorflow/
│ │ │ │ └── DISTS_tf.py
│ │ │ ├── LICENSE
│ │ │ ├── requirements.txt
│ │ │ └── weights/
│ │ │ └── alpha_beta.mat
│ │ ├── LPIPS.py
│ │ ├── PSNR.py
│ │ ├── dists.py
│ │ └── pip.sh
│ ├── README.md
│ ├── VERSION
│ ├── docs/
│ │ ├── CONTRIBUTING.md
│ │ ├── FAQ.md
│ │ ├── Training.md
│ │ ├── Training_CN.md
│ │ ├── anime_comparisons.md
│ │ ├── anime_comparisons_CN.md
│ │ ├── anime_model.md
│ │ ├── anime_video_model.md
│ │ ├── feedback.md
│ │ ├── model_zoo.md
│ │ └── ncnn_conversion.md
│ ├── kdsrgan/
│ │ ├── __init__.py
│ │ ├── archs/
│ │ │ ├── ST_arch.py
│ │ │ ├── TA_arch.py
│ │ │ ├── __init__.py
│ │ │ ├── common.py
│ │ │ ├── discriminator_arch.py
│ │ │ └── srvgg_arch.py
│ │ ├── data/
│ │ │ ├── __init__.py
│ │ │ ├── kdsrgan_dataset.py
│ │ │ └── kdsrgan_paired_dataset.py
│ │ ├── losses/
│ │ │ ├── __init__.py
│ │ │ └── my_loss.py
│ │ ├── models/
│ │ │ ├── __init__.py
│ │ │ ├── kdsrgan_ST_model.py
│ │ │ ├── kdsrgan_TA_model.py
│ │ │ ├── kdsrnet_ST_model.py
│ │ │ └── kdsrnet_TA_model.py
│ │ ├── test.py
│ │ ├── train.py
│ │ ├── utils.py
│ │ └── weights/
│ │ └── README.md
│ ├── options/
│ │ ├── test_kdsrgan_x4ST.yml
│ │ ├── test_kdsrnet_x4TA.yml
│ │ ├── train_kdsrgan_x4ST.yml
│ │ ├── train_kdsrgan_x4STV2.yml
│ │ ├── train_kdsrnet_x4ST.yml
│ │ └── train_kdsrnet_x4TA.yml
│ ├── pip.sh
│ ├── requirements.txt
│ ├── scripts/
│ │ ├── extract_subimages.py
│ │ ├── extract_subimages_DF2K.py
│ │ ├── generate_meta_info.py
│ │ ├── generate_meta_info_DF2K.py
│ │ ├── generate_meta_info_OST.py
│ │ ├── generate_meta_info_pairdata.py
│ │ ├── generate_multiscale_DF2K.py
│ │ └── pytorch2onnx.py
│ ├── setup.cfg
│ ├── setup.py
│ ├── test.sh
│ ├── tests/
│ │ ├── data/
│ │ │ ├── gt.lmdb/
│ │ │ │ ├── data.mdb
│ │ │ │ ├── lock.mdb
│ │ │ │ └── meta_info.txt
│ │ │ ├── lq.lmdb/
│ │ │ │ ├── data.mdb
│ │ │ │ ├── lock.mdb
│ │ │ │ └── meta_info.txt
│ │ │ ├── meta_info_gt.txt
│ │ │ ├── meta_info_pair.txt
│ │ │ ├── test_realesrgan_dataset.yml
│ │ │ ├── test_realesrgan_model.yml
│ │ │ ├── test_realesrgan_paired_dataset.yml
│ │ │ └── test_realesrnet_model.yml
│ │ ├── test_dataset.py
│ │ ├── test_discriminator_arch.py
│ │ ├── test_model.py
│ │ └── test_utils.py
│ ├── train_KDSRGAN_S.sh
│ ├── train_KDSRNet_S.sh
│ └── train_KDSRNet_T.sh
├── KDSR-classic/
│ ├── README.md
│ ├── cal_parms.py
│ ├── data/
│ │ ├── __init__.py
│ │ ├── benchmark.py
│ │ ├── common.py
│ │ ├── common2.py
│ │ ├── df2k.py
│ │ ├── multiscalesrdata.py
│ │ └── multiscalesrdata2.py
│ ├── dataloader.py
│ ├── loss/
│ │ ├── __init__.py
│ │ ├── adversarial.py
│ │ ├── discriminator.py
│ │ └── vgg.py
│ ├── main_anisonoise_KDSRsMx4_stage3.sh
│ ├── main_anisonoise_KDSRsMx4_stage4.sh
│ ├── main_anisonoise_stage3.py
│ ├── main_anisonoise_stage4.py
│ ├── main_iso_KDSRsLx4_stage3.sh
│ ├── main_iso_KDSRsLx4_stage4.sh
│ ├── main_iso_KDSRsMx4_stage3.sh
│ ├── main_iso_KDSRsMx4_stage4.sh
│ ├── main_iso_stage3.py
│ ├── main_iso_stage4.py
│ ├── model/
│ │ ├── __init__.py
│ │ ├── blindsr.py
│ │ └── common.py
│ ├── model_ST/
│ │ ├── __init__.py
│ │ ├── blindsr.py
│ │ └── common.py
│ ├── model_TA/
│ │ ├── __init__.py
│ │ ├── blindsr.py
│ │ └── common.py
│ ├── option.py
│ ├── template.py
│ ├── test.py
│ ├── test_anisonoise_KDSRsMx4.sh
│ ├── test_iso_KDSRsLx4.sh
│ ├── test_iso_KDSRsMx4.sh
│ ├── trainer_anisonoise_stage3.py
│ ├── trainer_anisonoise_stage4.py
│ ├── trainer_iso_stage3.py
│ ├── trainer_iso_stage4.py
│ ├── utility.py
│ ├── utility2.py
│ ├── utility3.py
│ └── utils/
│ ├── __init__.py
│ ├── util.py
│ └── util2.py
└── README.md
================================================
FILE CONTENTS
================================================
================================================
FILE: KDSR-GAN/.gitignore
================================================
# ignored folders
datasets/*
experiments/*
results/*
tb_logger/*
wandb/*
tmp/*
realesrgan/weights/*
version.py
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
pip-wheel-metadata/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
.python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
================================================
FILE: KDSR-GAN/.pre-commit-config.yaml
================================================
repos:
# flake8
- repo: https://github.com/PyCQA/flake8
rev: 3.8.3
hooks:
- id: flake8
args: ["--config=setup.cfg", "--ignore=W504, W503"]
# modify known_third_party
- repo: https://github.com/asottile/seed-isort-config
rev: v2.2.0
hooks:
- id: seed-isort-config
# isort
- repo: https://github.com/timothycrosley/isort
rev: 5.2.2
hooks:
- id: isort
# yapf
- repo: https://github.com/pre-commit/mirrors-yapf
rev: v0.30.0
hooks:
- id: yapf
# codespell
- repo: https://github.com/codespell-project/codespell
rev: v2.1.0
hooks:
- id: codespell
# pre-commit-hooks
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v3.2.0
hooks:
- id: trailing-whitespace # Trim trailing whitespace
- id: check-yaml # Attempt to load all yaml files to verify syntax
- id: check-merge-conflict # Check for files that contain merge conflict strings
- id: double-quote-string-fixer # Replace double quoted strings with single quoted strings
- id: end-of-file-fixer # Make sure files end in a newline and only a newline
- id: requirements-txt-fixer # Sort entries in requirements.txt and remove incorrect entry for pkg-resources==0.0.0
- id: fix-encoding-pragma # Remove the coding pragma: # -*- coding: utf-8 -*-
args: ["--remove"]
- id: mixed-line-ending # Replace or check mixed line ending
args: ["--fix=lf"]
================================================
FILE: KDSR-GAN/CODE_OF_CONDUCT.md
================================================
# Contributor Covenant Code of Conduct
## Our Pledge
We as members, contributors, and leaders pledge to make participation in our
community a harassment-free experience for everyone, regardless of age, body
size, visible or invisible disability, ethnicity, sex characteristics, gender
identity and expression, level of experience, education, socio-economic status,
nationality, personal appearance, race, religion, or sexual identity
and orientation.
We pledge to act and interact in ways that contribute to an open, welcoming,
diverse, inclusive, and healthy community.
## Our Standards
Examples of behavior that contributes to a positive environment for our
community include:
* Demonstrating empathy and kindness toward other people
* Being respectful of differing opinions, viewpoints, and experiences
* Giving and gracefully accepting constructive feedback
* Accepting responsibility and apologizing to those affected by our mistakes,
and learning from the experience
* Focusing on what is best not just for us as individuals, but for the
overall community
Examples of unacceptable behavior include:
* The use of sexualized language or imagery, and sexual attention or
advances of any kind
* Trolling, insulting or derogatory comments, and personal or political attacks
* Public or private harassment
* Publishing others' private information, such as a physical or email
address, without their explicit permission
* Other conduct which could reasonably be considered inappropriate in a
professional setting
## Enforcement Responsibilities
Community leaders are responsible for clarifying and enforcing our standards of
acceptable behavior and will take appropriate and fair corrective action in
response to any behavior that they deem inappropriate, threatening, offensive,
or harmful.
Community leaders have the right and responsibility to remove, edit, or reject
comments, commits, code, wiki edits, issues, and other contributions that are
not aligned to this Code of Conduct, and will communicate reasons for moderation
decisions when appropriate.
## Scope
This Code of Conduct applies within all community spaces, and also applies when
an individual is officially representing the community in public spaces.
Examples of representing our community include using an official e-mail address,
posting via an official social media account, or acting as an appointed
representative at an online or offline event.
## Enforcement
Instances of abusive, harassing, or otherwise unacceptable behavior may be
reported to the community leaders responsible for enforcement at
xintao.wang@outlook.com or xintaowang@tencent.com.
All complaints will be reviewed and investigated promptly and fairly.
All community leaders are obligated to respect the privacy and security of the
reporter of any incident.
## Enforcement Guidelines
Community leaders will follow these Community Impact Guidelines in determining
the consequences for any action they deem in violation of this Code of Conduct:
### 1. Correction
**Community Impact**: Use of inappropriate language or other behavior deemed
unprofessional or unwelcome in the community.
**Consequence**: A private, written warning from community leaders, providing
clarity around the nature of the violation and an explanation of why the
behavior was inappropriate. A public apology may be requested.
### 2. Warning
**Community Impact**: A violation through a single incident or series
of actions.
**Consequence**: A warning with consequences for continued behavior. No
interaction with the people involved, including unsolicited interaction with
those enforcing the Code of Conduct, for a specified period of time. This
includes avoiding interactions in community spaces as well as external channels
like social media. Violating these terms may lead to a temporary or
permanent ban.
### 3. Temporary Ban
**Community Impact**: A serious violation of community standards, including
sustained inappropriate behavior.
**Consequence**: A temporary ban from any sort of interaction or public
communication with the community for a specified period of time. No public or
private interaction with the people involved, including unsolicited interaction
with those enforcing the Code of Conduct, is allowed during this period.
Violating these terms may lead to a permanent ban.
### 4. Permanent Ban
**Community Impact**: Demonstrating a pattern of violation of community
standards, including sustained inappropriate behavior, harassment of an
individual, or aggression toward or disparagement of classes of individuals.
**Consequence**: A permanent ban from any sort of public interaction within
the community.
## Attribution
This Code of Conduct is adapted from the [Contributor Covenant][homepage],
version 2.0, available at
https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
Community Impact Guidelines were inspired by [Mozilla's code of conduct
enforcement ladder](https://github.com/mozilla/diversity).
[homepage]: https://www.contributor-covenant.org
For answers to common questions about this code of conduct, see the FAQ at
https://www.contributor-covenant.org/faq. Translations are available at
https://www.contributor-covenant.org/translations.
================================================
FILE: KDSR-GAN/Metric/DISTS/DISTS_pytorch/DISTS_pt.py
================================================
# This is a pytoch implementation of DISTS metric.
# Requirements: python >= 3.6, pytorch >= 1.0
import numpy as np
import os,sys
import torch
from torchvision import models,transforms
import torch.nn as nn
import torch.nn.functional as F
class L2pooling(nn.Module):
def __init__(self, filter_size=5, stride=2, channels=None, pad_off=0):
super(L2pooling, self).__init__()
self.padding = (filter_size - 2 )//2
self.stride = stride
self.channels = channels
a = np.hanning(filter_size)[1:-1]
g = torch.Tensor(a[:,None]*a[None,:])
g = g/torch.sum(g)
self.register_buffer('filter', g[None,None,:,:].repeat((self.channels,1,1,1)))
def forward(self, input):
input = input**2
out = F.conv2d(input, self.filter, stride=self.stride, padding=self.padding, groups=input.shape[1])
return (out+1e-12).sqrt()
class DISTS(torch.nn.Module):
def __init__(self, load_weights=True):
super(DISTS, self).__init__()
vgg_pretrained_features = models.vgg16(pretrained=True).features
self.stage1 = torch.nn.Sequential()
self.stage2 = torch.nn.Sequential()
self.stage3 = torch.nn.Sequential()
self.stage4 = torch.nn.Sequential()
self.stage5 = torch.nn.Sequential()
for x in range(0,4):
self.stage1.add_module(str(x), vgg_pretrained_features[x])
self.stage2.add_module(str(4), L2pooling(channels=64))
for x in range(5, 9):
self.stage2.add_module(str(x), vgg_pretrained_features[x])
self.stage3.add_module(str(9), L2pooling(channels=128))
for x in range(10, 16):
self.stage3.add_module(str(x), vgg_pretrained_features[x])
self.stage4.add_module(str(16), L2pooling(channels=256))
for x in range(17, 23):
self.stage4.add_module(str(x), vgg_pretrained_features[x])
self.stage5.add_module(str(23), L2pooling(channels=512))
for x in range(24, 30):
self.stage5.add_module(str(x), vgg_pretrained_features[x])
for param in self.parameters():
param.requires_grad = False
self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1,-1,1,1))
self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225]).view(1,-1,1,1))
self.chns = [3,64,128,256,512,512]
self.register_parameter("alpha", nn.Parameter(torch.randn(1, sum(self.chns),1,1)))
self.register_parameter("beta", nn.Parameter(torch.randn(1, sum(self.chns),1,1)))
self.alpha.data.normal_(0.1,0.01)
self.beta.data.normal_(0.1,0.01)
if load_weights:
# weights = torch.load(os.path.join(sys.prefix, 'weights.pt'))
weights = torch.load('scripts/metrics/DISTS/DISTS_pytorch/weights.pt')
self.alpha.data = weights['alpha']
self.beta.data = weights['beta']
def forward_once(self, x):
h = (x-self.mean)/self.std
h = self.stage1(h)
h_relu1_2 = h
h = self.stage2(h)
h_relu2_2 = h
h = self.stage3(h)
h_relu3_3 = h
h = self.stage4(h)
h_relu4_3 = h
h = self.stage5(h)
h_relu5_3 = h
return [x,h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3]
def forward(self, x, y, require_grad=False, batch_average=False):
if require_grad:
feats0 = self.forward_once(x)
feats1 = self.forward_once(y)
else:
with torch.no_grad():
feats0 = self.forward_once(x)
feats1 = self.forward_once(y)
dist1 = 0
dist2 = 0
c1 = 1e-6
c2 = 1e-6
w_sum = self.alpha.sum() + self.beta.sum()
alpha = torch.split(self.alpha/w_sum, self.chns, dim=1)
beta = torch.split(self.beta/w_sum, self.chns, dim=1)
for k in range(len(self.chns)):
x_mean = feats0[k].mean([2,3], keepdim=True)
y_mean = feats1[k].mean([2,3], keepdim=True)
S1 = (2*x_mean*y_mean+c1)/(x_mean**2+y_mean**2+c1)
dist1 = dist1+(alpha[k]*S1).sum(1,keepdim=True)
x_var = ((feats0[k]-x_mean)**2).mean([2,3], keepdim=True)
y_var = ((feats1[k]-y_mean)**2).mean([2,3], keepdim=True)
xy_cov = (feats0[k]*feats1[k]).mean([2,3],keepdim=True) - x_mean*y_mean
S2 = (2*xy_cov+c2)/(x_var+y_var+c2)
dist2 = dist2+(beta[k]*S2).sum(1,keepdim=True)
score = 1 - (dist1+dist2).squeeze()
if batch_average:
return score.mean()
else:
return score
def prepare_image(image, resize=True):
if resize and min(image.size) > 256:
image = transforms.functional.resize(image, 256)
image = transforms.ToTensor()(image)
return image.unsqueeze(0)
if __name__ == '__main__':
from PIL import Image
import glob
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
# others
# data_root = '/data1/liangjie/BasicSR_ALL/results/'
# ref_root = '/data1/liangjie/BasicSR_ALL/datasets/'
# ref_dirs = ['SISR_Test_matlab/Set5mod12', 'SISR_Test_matlab/Set14mod12', 'SISR_Test_matlab/Manga109mod12', 'SISR_Test_matlab/BSDS100mod12', 'SISR_Test_matlab/General100mod12', 'SISR_Test/Urban100', 'DIV2K/DIV2K_valid_HR/']
# datasets = ['Set5', 'Set14', 'Manga109', 'BSDS100', 'General100', 'Urban100', 'DIV2K100']
# img_dirs = ['SRGAN_official', 'ESRGAN_official', 'NatSR_official', 'USRGAN_official', 'SPSR_official', 'SPSR_DF2K', 'ESRGAN_ours_DIV2K', 'ESRGAN_ours_DIV2K_ema', 'ESRGAN_ours_DF2K', 'ESRGAN_ours_DF2K_ema']
# SFTGAN
# data_root = '/data1/liangjie/BasicSR_ALL/results'
# ref_root = '/data1/liangjie/BasicSR_ALL/results/SFTGAN_official'
# ref_dirs = ['GT'] * 7
# datasets = ['Set5', 'Set14', 'Manga109', 'BSDS100', 'General100', 'Urban100', 'DIV2K100']
# img_dirs = ['SFTGAN_official']
# new
data_root = 'results/'
ref_root = 'datasets/'
ref_dirs = ['DIV2K/DIV2K_valid_HR/']
datasets = ['DIV2K100']
img_dirs = ['ESRGAN_ours_DISTS_300k/visualization/']
logoverall_path = 'results/table_logs/' + 'DISTS_orisize_DISTStrain225k.txt'
for index in range(len(ref_dirs)):
ref_dir = os.path.join(ref_root, ref_dirs[index])
for method in img_dirs:
img_dir = os.path.join(data_root, method, datasets[index])
img_list = sorted(glob.glob(os.path.join(img_dir, '*')))
log_path = 'results/table_logs/' + img_dir.replace('/', '_') + '_DISTS_orisize.txt'
DISTS_all = []
for i, img_path in enumerate(img_list):
file_name = img_path.split('/')[-1]
if 'DIV2K100' in img_dir and 'SFTGAN' not in img_dir:
gt_path = os.path.join(ref_dir, file_name[:4] + '.png')
elif 'Urban100' in img_dir and 'SFTGAN' not in img_dir:
gt_path = os.path.join(ref_dir, file_name[:7] + '.png')
elif 'SFTGAN' in img_dir:
gt_path = os.path.join(ref_dir, file_name.split('_')[0] + '_gt.png')
if 'Urban100' in img_dir:
gt_path = os.path.join(ref_dir, file_name.split('_')[0] + '_' + file_name.split('_')[1] + '_gt.png')
else:
if '_' in file_name:
gt_path = os.path.join(ref_dir, file_name.split('_')[0] + '.png')
else:
gt_path = os.path.join(ref_dir, file_name)
ref = prepare_image(Image.open(gt_path).convert("RGB"), resize=False)
dist = prepare_image(Image.open(img_path).convert("RGB"), resize=False)
assert ref.shape == dist.shape
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = DISTS().to(device)
ref = ref.to(device)
dist = dist.to(device)
score = model(ref, dist)
DISTS_all.append(score.item())
log = f'{i + 1:3d}: {file_name:25}. \tDISTS: {score.item():.6f}.'
with open(log_path, 'a') as f:
f.write(log + '\n')
# print(log)
log = f'Average: DISTS: {sum(DISTS_all) / len(DISTS_all):.6f}'
with open(log_path, 'a') as f:
f.write(log + '\n')
log_overall = method + '__' + datasets[index] + '__' + log
with open(logoverall_path, 'a') as f:
f.write(log_overall + '\n')
print(log_overall)
================================================
FILE: KDSR-GAN/Metric/DISTS/DISTS_tensorflow/DISTS_tf.py
================================================
# This is a tensorflow implementation of DISTS metric.
# Requirements: python >= 3.6, tensorflow-gpu >= 1.15
import tensorflow.compat.v1 as tf
import numpy as np
import time
import scipy.io as scio
from PIL import Image
import argparse
# tf.enable_eager_execution()
tf.disable_eager_execution()
class DISTS():
def __init__(self):
self.parameters = scio.loadmat('../weights/net_param.mat')
self.chns = [3,64,128,256,512,512]
self.mean = tf.constant(self.parameters['vgg_mean'], dtype=tf.float32, shape=(1,1,1,3),name="img_mean")
self.std = tf.constant(self.parameters['vgg_std'], dtype=tf.float32, shape=(1,1,1,3),name="img_std")
# self.alpha = tf.Variable(tf.random_normal(shape=(1,1,1,sum(self.chns)), mean=0.1, stddev=0.01),name="alpha")
# self.beta = tf.Variable(tf.random_normal(shape=(1,1,1,sum(self.chns)), mean=0.1, stddev=0.01),name="beta")
self.weights = scio.loadmat('../weights/alpha_beta.mat')
self.alpha = tf.constant(np.reshape(self.weights['alpha'],(1,1,1,sum(self.chns))),name="alpha")
self.beta = tf.constant(np.reshape(self.weights['beta'],(1,1,1,sum(self.chns))),name="beta")
def get_features(self, img):
x = (img - self.mean)/self.std
self.conv1_1 = self.conv_layer(x, "conv1_1")
self.conv1_2 = self.conv_layer(self.conv1_1, "conv1_2")
self.pool1 = self.pool_layer(self.conv1_2, name="pool_1")
self.conv2_1 = self.conv_layer(self.pool1, "conv2_1")
self.conv2_2 = self.conv_layer(self.conv2_1, "conv2_2")
self.pool2 = self.pool_layer(self.conv2_2, name="pool_2")
self.conv3_1 = self.conv_layer(self.pool2, "conv3_1")
self.conv3_2 = self.conv_layer(self.conv3_1, "conv3_2")
self.conv3_3 = self.conv_layer(self.conv3_2, "conv3_3")
self.pool3 = self.pool_layer(self.conv3_3, name="pool_3")
self.conv4_1 = self.conv_layer(self.pool3, "conv4_1")
self.conv4_2 = self.conv_layer(self.conv4_1, "conv4_2")
self.conv4_3 = self.conv_layer(self.conv4_2, "conv4_3")
self.pool4 = self.pool_layer(self.conv4_3, name="pool_4")
self.conv5_1 = self.conv_layer(self.pool4, "conv5_1")
self.conv5_2 = self.conv_layer(self.conv5_1, "conv5_2")
self.conv5_3 = self.conv_layer(self.conv5_2, "conv5_3")
return [img, self.conv1_2,self.conv2_2,self.conv3_3,self.conv4_3,self.conv5_3]
def conv_layer(self, input, name):
with tf.variable_scope(name) as _:
filter = self.get_conv_filter(name)
conv = tf.nn.conv2d(input, filter, strides=1, padding="SAME")
bias = self.get_bias(name)
conv = tf.nn.relu(tf.nn.bias_add(conv, bias))
return conv
def pool_layer(self, input, name):
# return tf.nn.max_pool(input, ksize=[1,2,2,1], strides=[1,2,2,1], padding="SAME")
with tf.variable_scope(name) as _:
filter = tf.squeeze(tf.constant(self.parameters['L2'+name], name = "filter"),3)
conv = tf.nn.conv2d(input**2, filter, strides=2, padding=[[0, 0], [1, 0], [1, 0], [0, 0]])
return tf.sqrt(tf.maximum(conv, 1e-12))
def get_conv_filter(self, name):
return tf.constant(self.parameters[name+'_weight'], name = "filter")
def get_bias(self, name):
return tf.constant(np.squeeze(self.parameters[name+'_bias']), name = "bias")
def get_score(self, img1, img2):
feats0 = self.get_features(img1)
feats1 = self.get_features(img2)
dist1 = 0
dist2 = 0
c1 = 1e-6
c2 = 1e-6
w_sum = tf.reduce_sum(self.alpha) + tf.reduce_sum(self.beta)
alpha = tf.split(self.alpha/w_sum, self.chns, axis=3)
beta = tf.split(self.beta/w_sum, self.chns, axis=3)
for k in range(len(self.chns)):
x_mean = tf.reduce_mean(feats0[k],[1,2], keepdims=True)
y_mean = tf.reduce_mean(feats1[k],[1,2], keepdims=True)
S1 = (2*x_mean*y_mean+c1)/(x_mean**2+y_mean**2+c1)
dist1 = dist1+tf.reduce_sum(alpha[k]*S1, 3, keepdims=True)
x_var = tf.reduce_mean((feats0[k]-x_mean)**2,[1,2], keepdims=True)
y_var = tf.reduce_mean((feats1[k]-y_mean)**2,[1,2], keepdims=True)
xy_cov = tf.reduce_mean(feats0[k]*feats1[k],[1,2], keepdims=True) - x_mean*y_mean
S2 = (2*xy_cov+c2)/(x_var+y_var+c2)
dist2 = dist2+tf.reduce_sum(beta[k]*S2, 3, keepdims=True)
dist = 1-tf.squeeze(dist1+dist2)
return dist
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--ref', type=str, default='../images/r0.png')
parser.add_argument('--dist', type=str, default='../images/r1.png')
args = parser.parse_args()
model = DISTS()
ref = np.array(Image.open(args.ref).convert("RGB"))
ref = np.expand_dims(ref,axis=0)/255.
dist = np.array(Image.open(args.dist).convert("RGB"))
dist = np.expand_dims(dist,axis=0)/255.
x = tf.placeholder(dtype=tf.float32, shape=ref.shape, name= "ref")
y = tf.placeholder(dtype=tf.float32, shape=dist.shape, name= "dist")
score = model.get_score(x,y)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
score = sess.run(score, feed_dict={x: ref, y: dist})
print(score)
================================================
FILE: KDSR-GAN/Metric/DISTS/LICENSE
================================================
MIT License
Copyright (c) 2020 Keyan Ding
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
================================================
FILE: KDSR-GAN/Metric/DISTS/requirements.txt
================================================
torch>=1.0
================================================
FILE: KDSR-GAN/Metric/LPIPS.py
================================================
import cv2
import glob
import numpy as np
import os.path as osp
from torchvision.transforms.functional import normalize
from basicsr.utils import img2tensor
import lpips
import argparse
def main():
# Configurations
parser = argparse.ArgumentParser()
parser.add_argument('--folder_gt', type=str, default='/root/datasets/NTIRE2020-Track1/track1-valid-gt')
parser.add_argument('--folder_restored', type=str, default='/root/results/NTIRE2020-Track1')
args = parser.parse_args()
loss_fn_vgg = lpips.LPIPS(net='vgg').cuda(0)
lpips_all = []
img_list = sorted(glob.glob(osp.join(args.folder_gt, '*.png')))
lr_list = sorted(glob.glob(osp.join(args.folder_restored, '*.png')))
mean = [0.5, 0.5, 0.5]
std = [0.5, 0.5, 0.5]
for i, (img_path, lr_path) in enumerate(zip(img_list,lr_list)):
basename, ext = osp.splitext(osp.basename(img_path))
img_gt = cv2.imread(img_path, cv2.IMREAD_UNCHANGED).astype(np.float32) / 255.
img_restored = cv2.imread(osp.join(lr_path), cv2.IMREAD_UNCHANGED).astype(
np.float32) / 255.
img_gt, img_restored = img2tensor([img_gt, img_restored], bgr2rgb=True, float32=True)
# norm to [-1, 1]
normalize(img_gt, mean, std, inplace=True)
normalize(img_restored, mean, std, inplace=True)
# calculate lpips
lpips_val = loss_fn_vgg(img_restored.unsqueeze(0).cuda(0), img_gt.unsqueeze(0).cuda(0)).cpu().data.numpy()[0,0,0,0]
# print(lpips_val)
lpips_all.append(lpips_val)
print(f'Average: LPIPS: {sum(lpips_all) / len(lpips_all):.6f}')
if __name__ == '__main__':
main()
================================================
FILE: KDSR-GAN/Metric/PSNR.py
================================================
import cv2
import glob
import numpy as np
import os.path as osp
from torchvision.transforms.functional import normalize
from basicsr.utils import img2tensor
import lpips
import argparse
from basicsr.metrics import calculate_psnr, calculate_ssim
def main():
# Configurations
parser = argparse.ArgumentParser()
parser.add_argument('--folder_gt', type=str, default='/root/datasets/NTIRE2020-Track1/track1-valid-gt')
parser.add_argument('--folder_restored', type=str, default='/root/results/NTIRE2020-Track1')
args = parser.parse_args()
psnr_all = []
ssim_all = []
img_list = sorted(glob.glob(osp.join(args.folder_gt, '*.png')))
lr_list = sorted(glob.glob(osp.join(args.folder_restored, '*.png')))
for i, (img_path, lr_path) in enumerate(zip(img_list,lr_list)):
basename, ext = osp.splitext(osp.basename(img_path))
img_gt = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
img_restored = cv2.imread(osp.join(lr_path), cv2.IMREAD_UNCHANGED)
psnr=calculate_psnr(img_restored, img_gt, crop_border=4, test_y_channel=True)
ssim=calculate_ssim(img_restored, img_gt, crop_border=4, test_y_channel=True)
psnr_all.append(psnr)
ssim_all.append(ssim)
print(f'Average: PSNR: {sum(psnr_all) / len(psnr_all):.6f}')
print(f'Average: SSIM: {sum(ssim_all) / len(ssim_all):.6f}')
if __name__ == '__main__':
main()
================================================
FILE: KDSR-GAN/Metric/dists.py
================================================
# This is a pytoch implementation of DISTS metric.
# Requirements: python >= 3.6, pytorch >= 1.0
import numpy as np
import os,sys
import torch
from torchvision import models,transforms
import torch.nn as nn
import torch.nn.functional as F
import argparse
import os.path as osp
class L2pooling(nn.Module):
def __init__(self, filter_size=5, stride=2, channels=None, pad_off=0):
super(L2pooling, self).__init__()
self.padding = (filter_size - 2 )//2
self.stride = stride
self.channels = channels
a = np.hanning(filter_size)[1:-1]
g = torch.Tensor(a[:,None]*a[None,:])
g = g/torch.sum(g)
self.register_buffer('filter', g[None,None,:,:].repeat((self.channels,1,1,1)))
def forward(self, input):
input = input**2
out = F.conv2d(input, self.filter, stride=self.stride, padding=self.padding, groups=input.shape[1])
return (out+1e-12).sqrt()
class DISTS(torch.nn.Module):
def __init__(self, load_weights=True):
super(DISTS, self).__init__()
vgg_pretrained_features = models.vgg16(pretrained=True).features
self.stage1 = torch.nn.Sequential()
self.stage2 = torch.nn.Sequential()
self.stage3 = torch.nn.Sequential()
self.stage4 = torch.nn.Sequential()
self.stage5 = torch.nn.Sequential()
for x in range(0,4):
self.stage1.add_module(str(x), vgg_pretrained_features[x])
self.stage2.add_module(str(4), L2pooling(channels=64))
for x in range(5, 9):
self.stage2.add_module(str(x), vgg_pretrained_features[x])
self.stage3.add_module(str(9), L2pooling(channels=128))
for x in range(10, 16):
self.stage3.add_module(str(x), vgg_pretrained_features[x])
self.stage4.add_module(str(16), L2pooling(channels=256))
for x in range(17, 23):
self.stage4.add_module(str(x), vgg_pretrained_features[x])
self.stage5.add_module(str(23), L2pooling(channels=512))
for x in range(24, 30):
self.stage5.add_module(str(x), vgg_pretrained_features[x])
for param in self.parameters():
param.requires_grad = False
self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1,-1,1,1))
self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225]).view(1,-1,1,1))
self.chns = [3,64,128,256,512,512]
self.register_parameter("alpha", nn.Parameter(torch.randn(1, sum(self.chns),1,1)))
self.register_parameter("beta", nn.Parameter(torch.randn(1, sum(self.chns),1,1)))
self.alpha.data.normal_(0.1,0.01)
self.beta.data.normal_(0.1,0.01)
if load_weights:
# weights = torch.load(os.path.join(sys.prefix, 'weights.pt'))
weights = torch.load('DISTS/DISTS_pytorch/weights.pt')
self.alpha.data = weights['alpha']
self.beta.data = weights['beta']
def forward_once(self, x):
h = (x-self.mean)/self.std
h = self.stage1(h)
h_relu1_2 = h
h = self.stage2(h)
h_relu2_2 = h
h = self.stage3(h)
h_relu3_3 = h
h = self.stage4(h)
h_relu4_3 = h
h = self.stage5(h)
h_relu5_3 = h
return [x,h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3]
def forward(self, x, y, require_grad=False, batch_average=False):
if require_grad:
feats0 = self.forward_once(x)
feats1 = self.forward_once(y)
else:
with torch.no_grad():
feats0 = self.forward_once(x)
feats1 = self.forward_once(y)
dist1 = 0
dist2 = 0
c1 = 1e-6
c2 = 1e-6
w_sum = self.alpha.sum() + self.beta.sum()
alpha = torch.split(self.alpha/w_sum, self.chns, dim=1)
beta = torch.split(self.beta/w_sum, self.chns, dim=1)
for k in range(len(self.chns)):
x_mean = feats0[k].mean([2,3], keepdim=True)
y_mean = feats1[k].mean([2,3], keepdim=True)
S1 = (2*x_mean*y_mean+c1)/(x_mean**2+y_mean**2+c1)
dist1 = dist1+(alpha[k]*S1).sum(1,keepdim=True)
x_var = ((feats0[k]-x_mean)**2).mean([2,3], keepdim=True)
y_var = ((feats1[k]-y_mean)**2).mean([2,3], keepdim=True)
xy_cov = (feats0[k]*feats1[k]).mean([2,3],keepdim=True) - x_mean*y_mean
S2 = (2*xy_cov+c2)/(x_var+y_var+c2)
dist2 = dist2+(beta[k]*S2).sum(1,keepdim=True)
score = 1 - (dist1+dist2).squeeze()
if batch_average:
return score.mean()
else:
return score
def prepare_image(image, resize=True):
if resize and min(image.size) > 256:
image = transforms.functional.resize(image, 256)
image = transforms.ToTensor()(image)
return image.unsqueeze(0)
if __name__ == '__main__':
from PIL import Image
import glob
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
parser = argparse.ArgumentParser()
parser.add_argument('--folder_gt', type=str, default='/root/datasets/NTIRE2020-Track1/track1-valid-gt')
parser.add_argument('--folder_restored', type=str, default='/root/results/NTIRE2020-Track1')
args = parser.parse_args()
img_list = sorted(glob.glob(osp.join(args.folder_gt, '*.png')))
lr_list = sorted(glob.glob(osp.join(args.folder_restored, '*.png')))
DISTS_all = []
for i, (gt_path, lr_path) in enumerate(zip(img_list,lr_list)):
ref = prepare_image(Image.open(gt_path).convert("RGB"), resize=False)
dist = prepare_image(Image.open(lr_path).convert("RGB"), resize=False)
assert ref.shape == dist.shape
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = DISTS().to(device)
ref = ref.to(device)
dist = dist.to(device)
score = model(ref, dist)
DISTS_all.append(score.item())
log = f'{i + 1:3d}:\tDISTS: {score.item():.6f}.'
print(log)
# with open(log_path, 'a') as f:
# f.write(log + '\n')
# # print(log)
print(f'Average: DISTS: {sum(DISTS_all) / len(DISTS_all):.6f}')
================================================
FILE: KDSR-GAN/Metric/pip.sh
================================================
pip install lpips --index-url http://pypi.douban.com/simple --trusted-host pypi.douban.com
pip install basicsr --index-url http://pypi.douban.com/simple --trusted-host pypi.douban.com
================================================
FILE: KDSR-GAN/README.md
================================================
# KDSR-GAN
This project is the official implementation of 'Knowledge Distillation based Degradation Estimation for Blind Super-Resolution', ICLR2023
> **Knowledge Distillation based Degradation Estimation for Blind Super-Resolution [[Paper](https://arxiv.org/pdf/2211.16928.pdf)] [[Project](https://github.com/Zj-BinXia/KDSR)]**
This is code for KDSR-GAN (for Real-world SR)
<p align="center">
<img src="images/method.jpg" width="50%">
</p>
---
## Dependencies and Installation
- Python >= 3.8 (Recommend to use [Anaconda](https://www.anaconda.com/download/#linux) or [Miniconda](https://docs.conda.io/en/latest/miniconda.html))
- [PyTorch >= 1.10](https://pytorch.org/)
### Installation
Install dependent packages
pip install basicsr
pip install -r requirements.txt
pip install pandas
sudo python3 setup.py develop
---
## Dataset Preparation
We use the same training datasets as [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN) (DF2K+OST).
---
## Training (8 V100 GPUs)
1. We train KDSRNet_T (only using L1 loss)
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python3 -m torch.distributed.launch --nproc_per_node=8 --master_port=3309 kdsrgan/train.py -opt options/train_kdsrnet_x4TA.yml --launcher pytorch
```
2. we train KDSRNet_S (only using L1 loss and KD loss). **It is notable that modify the ''pretrain_network_TA'' and ''pretrain_network_g'' of options/train_kdsrnet_x4ST.yml to the path of trained KDSRNet_T checkpoint.** Then, we run
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
python3 -m torch.distributed.launch --nproc_per_node=8 --master_port=4349 kdsrgan/train.py -opt options/train_kdsrnet_x4ST.yml --launcher pytorch
```
3. we train KDSRGAN_S ( using L1 loss, perceptual loss, adversial loss and KD loss). **It is notable that modify the ''pretrain_network_TA'' and ''pretrain_network_g'' of options/train_kdsrnet_x4ST.yml or options/train_kdsrnet_x4STV2.yml to the path of trained KDSRNet_T and KDSRNet_S checkpoint, respectively.** Then, we run
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
python3 -m torch.distributed.launch --nproc_per_node=8 --master_port=4397 kdsrgan/train.py -opt options/train_kdsrgan_x4ST.yml --launcher pytorch
```
or
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
python3 -m torch.distributed.launch --nproc_per_node=8 --master_port=4397 kdsrgan/train.py -opt options/train_kdsrgan_x4STV2.yml --launcher pytorch
```
---
## :european_castle: Model Zoo
Please download checkpoints from [Google Drive](https://drive.google.com/drive/folders/1QlOz4F9Mtp9DFXoaHYbnMnRSonR9YFJA).
---
## Testing
You can download NTIRE2020 Track1 and NTIRE2020 Track2 and AIM2019 Track2 from [official link](https://competitions.codalab.org/competitions/22220) or [Google Drive](https://drive.google.com/drive/folders/1P5x8wuty3IVj7aKoDe6uZy9pA19S9rrQ?usp=sharing).
```bash
python3 kdsrgan/test.py -opt options/test_kdsrgan_x4ST.yml
```
calculate metric
```bash
python3 Metric/PSNR.py --folder_gt PathtoGT --folder_restored PathtoSR
```
```bash
python3 Metric/LPIPS.py --folder_gt PathtoGT --folder_restored PathtoSR
```
---
## Results
<p align="center">
<img src="images/quan.jpg" width="90%">
</p>
---
## BibTeX
@InProceedings{xia2022knowledge,
title={Knowledge Distillation based Degradation Estimation for Blind Super-Resolution},
author={Xia, Bin and Zhang, Yulun and Wang, Yitong and Tian, Yapeng and Yang, Wenming and Timofte, Radu and Van Gool, Luc},
journal={ICLR},
year={2023}
}
## 📧 Contact
If you have any question, please email `zjbinxia@gmail.com`.
================================================
FILE: KDSR-GAN/VERSION
================================================
0.2.5.0
================================================
FILE: KDSR-GAN/docs/CONTRIBUTING.md
================================================
# Contributing to Real-ESRGAN
:art: Real-ESRGAN needs your contributions. Any contributions are welcome, such as new features/models/typo fixes/suggestions/maintenance, *etc*. See [CONTRIBUTING.md](docs/CONTRIBUTING.md). All contributors are list [here](README.md#hugs-acknowledgement).
We like open-source and want to develop practical algorithms for general image restoration. However, individual strength is limited. So, any kinds of contributions are welcome, such as:
- New features
- New models (your fine-tuned models)
- Bug fixes
- Typo fixes
- Suggestions
- Maintenance
- Documents
- *etc*
## Workflow
1. Fork and pull the latest Real-ESRGAN repository
1. Checkout a new branch (do not use master branch for PRs)
1. Commit your changes
1. Create a PR
**Note**:
1. Please check the code style and linting
1. The style configuration is specified in [setup.cfg](setup.cfg)
1. If you use VSCode, the settings are configured in [.vscode/settings.json](.vscode/settings.json)
1. Strongly recommend using `pre-commit hook`. It will check your code style and linting before your commit.
1. In the root path of project folder, run `pre-commit install`
1. The pre-commit configuration is listed in [.pre-commit-config.yaml](.pre-commit-config.yaml)
1. Better to [open a discussion](https://github.com/xinntao/Real-ESRGAN/discussions) before large changes.
1. Welcome to discuss :sunglasses:. I will try my best to join the discussion.
## TODO List
:zero: The most straightforward way of improving model performance is to fine-tune on some specific datasets.
Here are some TODOs:
- [ ] optimize for human faces
- [ ] optimize for texts
- [ ] support controllable restoration strength
:one: There are also [several issues](https://github.com/xinntao/Real-ESRGAN/issues) that require helpers to improve. If you can help, please let me know :smile:
================================================
FILE: KDSR-GAN/docs/FAQ.md
================================================
# FAQ
1. **Q: How to select models?**<br>
A: Please refer to [docs/model_zoo.md](docs/model_zoo.md)
1. **Q: Can `face_enhance` be used for anime images/animation videos?**<br>
A: No, it can only be used for real faces. It is recommended not to use this option for anime images/animation videos to save GPU memory.
1. **Q: Error "slow_conv2d_cpu" not implemented for 'Half'**<br>
A: In order to save GPU memory consumption and speed up inference, Real-ESRGAN uses half precision (fp16) during inference by default. However, some operators for half inference are not implemented in CPU mode. You need to add **`--fp32` option** for the commands. For example, `python inference_realesrgan.py -n RealESRGAN_x4plus.pth -i inputs --fp32`.
================================================
FILE: KDSR-GAN/docs/Training.md
================================================
# :computer: How to Train/Finetune Real-ESRGAN
- [Train Real-ESRGAN](#train-real-esrgan)
- [Overview](#overview)
- [Dataset Preparation](#dataset-preparation)
- [Train Real-ESRNet](#Train-Real-ESRNet)
- [Train Real-ESRGAN](#Train-Real-ESRGAN)
- [Finetune Real-ESRGAN on your own dataset](#Finetune-Real-ESRGAN-on-your-own-dataset)
- [Generate degraded images on the fly](#Generate-degraded-images-on-the-fly)
- [Use paired training data](#use-your-own-paired-data)
[English](Training.md) **|** [简体中文](Training_CN.md)
## Train Real-ESRGAN
### Overview
The training has been divided into two stages. These two stages have the same data synthesis process and training pipeline, except for the loss functions. Specifically,
1. We first train Real-ESRNet with L1 loss from the pre-trained model ESRGAN.
1. We then use the trained Real-ESRNet model as an initialization of the generator, and train the Real-ESRGAN with a combination of L1 loss, perceptual loss and GAN loss.
### Dataset Preparation
We use DF2K (DIV2K and Flickr2K) + OST datasets for our training. Only HR images are required. <br>
You can download from :
1. DIV2K: http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_HR.zip
2. Flickr2K: https://cv.snu.ac.kr/research/EDSR/Flickr2K.tar
3. OST: https://openmmlab.oss-cn-hangzhou.aliyuncs.com/datasets/OST_dataset.zip
Here are steps for data preparation.
#### Step 1: [Optional] Generate multi-scale images
For the DF2K dataset, we use a multi-scale strategy, *i.e.*, we downsample HR images to obtain several Ground-Truth images with different scales. <br>
You can use the [scripts/generate_multiscale_DF2K.py](scripts/generate_multiscale_DF2K.py) script to generate multi-scale images. <br>
Note that this step can be omitted if you just want to have a fast try.
```bash
python scripts/generate_multiscale_DF2K.py --input datasets/DF2K/DF2K_HR --output datasets/DF2K/DF2K_multiscale
```
#### Step 2: [Optional] Crop to sub-images
We then crop DF2K images into sub-images for faster IO and processing.<br>
This step is optional if your IO is enough or your disk space is limited.
You can use the [scripts/extract_subimages.py](scripts/extract_subimages.py) script. Here is the example:
```bash
python scripts/extract_subimages.py --input datasets/DF2K/DF2K_multiscale --output datasets/DF2K/DF2K_multiscale_sub --crop_size 400 --step 200
```
#### Step 3: Prepare a txt for meta information
You need to prepare a txt file containing the image paths. The following are some examples in `meta_info_DF2Kmultiscale+OST_sub.txt` (As different users may have different sub-images partitions, this file is not suitable for your purpose and you need to prepare your own txt file):
```txt
DF2K_HR_sub/000001_s001.png
DF2K_HR_sub/000001_s002.png
DF2K_HR_sub/000001_s003.png
...
```
You can use the [scripts/generate_meta_info.py](scripts/generate_meta_info.py) script to generate the txt file. <br>
You can merge several folders into one meta_info txt. Here is the example:
```bash
python scripts/generate_meta_info.py --input datasets/DF2K/DF2K_HR datasets/DF2K/DF2K_multiscale --root datasets/DF2K datasets/DF2K --meta_info datasets/DF2K/meta_info/meta_info_DF2Kmultiscale.txt
```
### Train Real-ESRNet
1. Download pre-trained model [ESRGAN](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth) into `experiments/pretrained_models`.
```bash
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth -P experiments/pretrained_models
```
1. Modify the content in the option file `options/train_realesrnet_x4plus.yml` accordingly:
```yml
train:
name: DF2K+OST
type: RealESRGANDataset
dataroot_gt: datasets/DF2K # modify to the root path of your folder
meta_info: realesrgan/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt # modify to your own generate meta info txt
io_backend:
type: disk
```
1. If you want to perform validation during training, uncomment those lines and modify accordingly:
```yml
# Uncomment these for validation
# val:
# name: validation
# type: PairedImageDataset
# dataroot_gt: path_to_gt
# dataroot_lq: path_to_lq
# io_backend:
# type: disk
...
# Uncomment these for validation
# validation settings
# val:
# val_freq: !!float 5e3
# save_img: True
# metrics:
# psnr: # metric name, can be arbitrary
# type: calculate_psnr
# crop_border: 4
# test_y_channel: false
```
1. Before the formal training, you may run in the `--debug` mode to see whether everything is OK. We use four GPUs for training:
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/train_realesrnet_x4plus.yml --launcher pytorch --debug
```
Train with **a single GPU** in the *debug* mode:
```bash
python realesrgan/train.py -opt options/train_realesrnet_x4plus.yml --debug
```
1. The formal training. We use four GPUs for training. We use the `--auto_resume` argument to automatically resume the training if necessary.
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/train_realesrnet_x4plus.yml --launcher pytorch --auto_resume
```
Train with **a single GPU**:
```bash
python realesrgan/train.py -opt options/train_realesrnet_x4plus.yml --auto_resume
```
### Train Real-ESRGAN
1. After the training of Real-ESRNet, you now have the file `experiments/train_RealESRNetx4plus_1000k_B12G4_fromESRGAN/model/net_g_1000000.pth`. If you need to specify the pre-trained path to other files, modify the `pretrain_network_g` value in the option file `train_realesrgan_x4plus.yml`.
1. Modify the option file `train_realesrgan_x4plus.yml` accordingly. Most modifications are similar to those listed above.
1. Before the formal training, you may run in the `--debug` mode to see whether everything is OK. We use four GPUs for training:
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/train_realesrgan_x4plus.yml --launcher pytorch --debug
```
Train with **a single GPU** in the *debug* mode:
```bash
python realesrgan/train.py -opt options/train_realesrgan_x4plus.yml --debug
```
1. The formal training. We use four GPUs for training. We use the `--auto_resume` argument to automatically resume the training if necessary.
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/train_realesrgan_x4plus.yml --launcher pytorch --auto_resume
```
Train with **a single GPU**:
```bash
python realesrgan/train.py -opt options/train_realesrgan_x4plus.yml --auto_resume
```
## Finetune Real-ESRGAN on your own dataset
You can finetune Real-ESRGAN on your own dataset. Typically, the fine-tuning process can be divided into two cases:
1. [Generate degraded images on the fly](#Generate-degraded-images-on-the-fly)
1. [Use your own **paired** data](#Use-paired-training-data)
### Generate degraded images on the fly
Only high-resolution images are required. The low-quality images are generated with the degradation process described in Real-ESRGAN during trainig.
**1. Prepare dataset**
See [this section](#dataset-preparation) for more details.
**2. Download pre-trained models**
Download pre-trained models into `experiments/pretrained_models`.
- *RealESRGAN_x4plus.pth*:
```bash
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P experiments/pretrained_models
```
- *RealESRGAN_x4plus_netD.pth*:
```bash
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.3/RealESRGAN_x4plus_netD.pth -P experiments/pretrained_models
```
**3. Finetune**
Modify [options/finetune_realesrgan_x4plus.yml](options/finetune_realesrgan_x4plus.yml) accordingly, especially the `datasets` part:
```yml
train:
name: DF2K+OST
type: RealESRGANDataset
dataroot_gt: datasets/DF2K # modify to the root path of your folder
meta_info: realesrgan/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt # modify to your own generate meta info txt
io_backend:
type: disk
```
We use four GPUs for training. We use the `--auto_resume` argument to automatically resume the training if necessary.
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/finetune_realesrgan_x4plus.yml --launcher pytorch --auto_resume
```
Finetune with **a single GPU**:
```bash
python realesrgan/train.py -opt options/finetune_realesrgan_x4plus.yml --auto_resume
```
### Use your own paired data
You can also finetune RealESRGAN with your own paired data. It is more similar to fine-tuning ESRGAN.
**1. Prepare dataset**
Assume that you already have two folders:
- **gt folder** (Ground-truth, high-resolution images): *datasets/DF2K/DIV2K_train_HR_sub*
- **lq folder** (Low quality, low-resolution images): *datasets/DF2K/DIV2K_train_LR_bicubic_X4_sub*
Then, you can prepare the meta_info txt file using the script [scripts/generate_meta_info_pairdata.py](scripts/generate_meta_info_pairdata.py):
```bash
python scripts/generate_meta_info_pairdata.py --input datasets/DF2K/DIV2K_train_HR_sub datasets/DF2K/DIV2K_train_LR_bicubic_X4_sub --meta_info datasets/DF2K/meta_info/meta_info_DIV2K_sub_pair.txt
```
**2. Download pre-trained models**
Download pre-trained models into `experiments/pretrained_models`.
- *RealESRGAN_x4plus.pth*
```bash
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P experiments/pretrained_models
```
- *RealESRGAN_x4plus_netD.pth*
```bash
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.3/RealESRGAN_x4plus_netD.pth -P experiments/pretrained_models
```
**3. Finetune**
Modify [options/finetune_realesrgan_x4plus_pairdata.yml](options/finetune_realesrgan_x4plus_pairdata.yml) accordingly, especially the `datasets` part:
```yml
train:
name: DIV2K
type: RealESRGANPairedDataset
dataroot_gt: datasets/DF2K # modify to the root path of your folder
dataroot_lq: datasets/DF2K # modify to the root path of your folder
meta_info: datasets/DF2K/meta_info/meta_info_DIV2K_sub_pair.txt # modify to your own generate meta info txt
io_backend:
type: disk
```
We use four GPUs for training. We use the `--auto_resume` argument to automatically resume the training if necessary.
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/finetune_realesrgan_x4plus_pairdata.yml --launcher pytorch --auto_resume
```
Finetune with **a single GPU**:
```bash
python realesrgan/train.py -opt options/finetune_realesrgan_x4plus_pairdata.yml --auto_resume
```
================================================
FILE: KDSR-GAN/docs/Training_CN.md
================================================
# :computer: 如何训练/微调 Real-ESRGAN
- [训练 Real-ESRGAN](#训练-real-esrgan)
- [概述](#概述)
- [准备数据集](#准备数据集)
- [训练 Real-ESRNet 模型](#训练-real-esrnet-模型)
- [训练 Real-ESRGAN 模型](#训练-real-esrgan-模型)
- [用自己的数据集微调 Real-ESRGAN](#用自己的数据集微调-real-esrgan)
- [动态生成降级图像](#动态生成降级图像)
- [使用已配对的数据](#使用已配对的数据)
[English](Training.md) **|** [简体中文](Training_CN.md)
## 训练 Real-ESRGAN
### 概述
训练分为两个步骤。除了 loss 函数外，这两个步骤拥有相同数据合成以及训练的一条龙流程。具体点说：
1. 首先使用 L1 loss 训练 Real-ESRNet 模型，其中 L1 loss 来自预先训练的 ESRGAN 模型。
2. 然后我们将 Real-ESRNet 模型作为生成器初始化，结合L1 loss、感知 loss、GAN loss 三者的参数对 Real-ESRGAN 进行训练。
### 准备数据集
我们使用 DF2K ( DIV2K 和 Flickr2K ) + OST 数据集进行训练。只需要HR图像！<br>
下面是网站链接:
1. DIV2K: http://data.vision.ee.ethz.ch/cvl/DIV2K/DIV2K_train_HR.zip
2. Flickr2K: https://cv.snu.ac.kr/research/EDSR/Flickr2K.tar
3. OST: https://openmmlab.oss-cn-hangzhou.aliyuncs.com/datasets/OST_dataset.zip
以下是数据的准备步骤。
#### 第1步：【可选】生成多尺寸图片
针对 DF2K 数据集，我们使用多尺寸缩放策略，*换言之*，我们对 HR 图像进行下采样，就能获得多尺寸的标准参考（Ground-Truth）图像。 <br>
您可以使用这个 [scripts/generate_multiscale_DF2K.py](scripts/generate_multiscale_DF2K.py) 脚本快速生成多尺寸的图像。<br>
注意：如果您只想简单试试，那么可以跳过此步骤。
```bash
python scripts/generate_multiscale_DF2K.py --input datasets/DF2K/DF2K_HR --output datasets/DF2K/DF2K_multiscale
```
#### 第2步：【可选】裁切为子图像
我们可以将 DF2K 图像裁切为子图像，以加快 IO 和处理速度。<br>
如果你的 IO 够好或储存空间有限，那么此步骤是可选的。<br>
您可以使用脚本 [scripts/extract_subimages.py](scripts/extract_subimages.py)。这是使用示例:
```bash
python scripts/extract_subimages.py --input datasets/DF2K/DF2K_multiscale --output datasets/DF2K/DF2K_multiscale_sub --crop_size 400 --step 200
```
#### 第3步：准备元信息 txt
您需要准备一个包含图像路径的 txt 文件。下面是 `meta_info_DF2Kmultiscale+OST_sub.txt` 中的部分展示（由于各个用户可能有截然不同的子图像划分，这个文件不适合你的需求，你得准备自己的 txt 文件)：
```txt
DF2K_HR_sub/000001_s001.png
DF2K_HR_sub/000001_s002.png
DF2K_HR_sub/000001_s003.png
...
```
你可以使用该脚本 [scripts/generate_meta_info.py](scripts/generate_meta_info.py) 生成包含图像路径的 txt 文件。<br>
你还可以合并多个文件夹的图像路径到一个元信息（meta_info）txt。这是使用示例:
```bash
python scripts/generate_meta_info.py --input datasets/DF2K/DF2K_HR, datasets/DF2K/DF2K_multiscale --root datasets/DF2K, datasets/DF2K --meta_info datasets/DF2K/meta_info/meta_info_DF2Kmultiscale.txt
```
### 训练 Real-ESRNet 模型
1. 下载预先训练的模型 [ESRGAN](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth)，放到 `experiments/pretrained_models`目录下。
```bash
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth -P experiments/pretrained_models
```
2. 相应地修改选项文件 `options/train_realesrnet_x4plus.yml` 中的内容：
```yml
train:
name: DF2K+OST
type: RealESRGANDataset
dataroot_gt: datasets/DF2K # 修改为你的数据集文件夹根目录
meta_info: realesrgan/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt # 修改为你自己生成的元信息txt
io_backend:
type: disk
```
3. 如果你想在训练过程中执行验证，就取消注释这些内容并进行相应的修改：
```yml
# 取消注释这些以进行验证
# val:
# name: validation
# type: PairedImageDataset
# dataroot_gt: path_to_gt
# dataroot_lq: path_to_lq
# io_backend:
# type: disk
...
# 取消注释这些以进行验证
# 验证设置
# val:
# val_freq: !!float 5e3
# save_img: True
# metrics:
# psnr: # 指标名称，可以是任意的
# type: calculate_psnr
# crop_border: 4
# test_y_channel: false
```
4. 正式训练之前，你可以用 `--debug` 模式检查是否正常运行。我们用了4个GPU进行训练：
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/train_realesrnet_x4plus.yml --launcher pytorch --debug
```
用 **1个GPU** 训练的 debug 模式示例:
```bash
python realesrgan/train.py -opt options/train_realesrnet_x4plus.yml --debug
```
5. 正式训练开始。我们用了4个GPU进行训练。还可以使用参数 `--auto_resume` 在必要时自动恢复训练。
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/train_realesrnet_x4plus.yml --launcher pytorch --auto_resume
```
用 **1个GPU** 训练：
```bash
python realesrgan/train.py -opt options/train_realesrnet_x4plus.yml --auto_resume
```
### 训练 Real-ESRGAN 模型
1. 训练 Real-ESRNet 模型后，您得到了这个 `experiments/train_RealESRNetx4plus_1000k_B12G4_fromESRGAN/model/net_g_1000000.pth` 文件。如果需要指定预训练路径到其他文件，请修改选项文件 `train_realesrgan_x4plus.yml` 中 `pretrain_network_g` 的值。
1. 修改选项文件 `train_realesrgan_x4plus.yml` 的内容。大多数修改与上节提到的类似。
1. 正式训练之前，你可以以 `--debug` 模式检查是否正常运行。我们使用了4个GPU进行训练：
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/train_realesrgan_x4plus.yml --launcher pytorch --debug
```
用 **1个GPU** 训练的 debug 模式示例:
```bash
python realesrgan/train.py -opt options/train_realesrgan_x4plus.yml --debug
```
1. 正式训练开始。我们使用4个GPU进行训练。还可以使用参数 `--auto_resume` 在必要时自动恢复训练。
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/train_realesrgan_x4plus.yml --launcher pytorch --auto_resume
```
用 **1个GPU** 训练：
```bash
python realesrgan/train.py -opt options/train_realesrgan_x4plus.yml --auto_resume
```
## 用自己的数据集微调 Real-ESRGAN
你可以用自己的数据集微调 Real-ESRGAN。一般地，微调（Fine-Tune）程序可以分为两种类型:
1. [动态生成降级图像](#动态生成降级图像)
2. [使用**已配对**的数据](#使用已配对的数据)
### 动态生成降级图像
只需要高分辨率图像。在训练过程中，使用 Real-ESRGAN 描述的降级模型生成低质量图像。
**1. 准备数据集**
完整信息请参见[本节](#准备数据集)。
**2. 下载预训练模型**
下载预先训练的模型到 `experiments/pretrained_models` 目录下。
- *RealESRGAN_x4plus.pth*:
```bash
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P experiments/pretrained_models
```
- *RealESRGAN_x4plus_netD.pth*:
```bash
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.3/RealESRGAN_x4plus_netD.pth -P experiments/pretrained_models
```
**3. 微调**
修改选项文件 [options/finetune_realesrgan_x4plus.yml](options/finetune_realesrgan_x4plus.yml) ，特别是 `datasets` 部分：
```yml
train:
name: DF2K+OST
type: RealESRGANDataset
dataroot_gt: datasets/DF2K # 修改为你的数据集文件夹根目录
meta_info: realesrgan/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt # 修改为你自己生成的元信息txt
io_backend:
type: disk
```
我们使用4个GPU进行训练。还可以使用参数 `--auto_resume` 在必要时自动恢复训练。
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/finetune_realesrgan_x4plus.yml --launcher pytorch --auto_resume
```
用 **1个GPU** 训练：
```bash
python realesrgan/train.py -opt options/finetune_realesrgan_x4plus.yml --auto_resume
```
### 使用已配对的数据
你还可以用自己已经配对的数据微调 RealESRGAN。这个过程更类似于微调 ESRGAN。
**1. 准备数据集**
假设你已经有两个文件夹（folder）:
- **gt folder**（标准参考，高分辨率图像）：*datasets/DF2K/DIV2K_train_HR_sub*
- **lq folder**（低质量，低分辨率图像）：*datasets/DF2K/DIV2K_train_LR_bicubic_X4_sub*
然后，您可以使用脚本 [scripts/generate_meta_info_pairdata.py](scripts/generate_meta_info_pairdata.py) 生成元信息（meta_info）txt 文件。
```bash
python scripts/generate_meta_info_pairdata.py --input datasets/DF2K/DIV2K_train_HR_sub datasets/DF2K/DIV2K_train_LR_bicubic_X4_sub --meta_info datasets/DF2K/meta_info/meta_info_DIV2K_sub_pair.txt
```
**2. 下载预训练模型**
下载预先训练的模型到 `experiments/pretrained_models` 目录下。
- *RealESRGAN_x4plus.pth*:
```bash
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P experiments/pretrained_models
```
- *RealESRGAN_x4plus_netD.pth*:
```bash
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.3/RealESRGAN_x4plus_netD.pth -P experiments/pretrained_models
```
**3. 微调**
修改选项文件 [options/finetune_realesrgan_x4plus_pairdata.yml](options/finetune_realesrgan_x4plus_pairdata.yml) ，特别是 `datasets` 部分：
```yml
train:
name: DIV2K
type: RealESRGANPairedDataset
dataroot_gt: datasets/DF2K # 修改为你的 gt folder 文件夹根目录
dataroot_lq: datasets/DF2K # 修改为你的 lq folder 文件夹根目录
meta_info: datasets/DF2K/meta_info/meta_info_DIV2K_sub_pair.txt # 修改为你自己生成的元信息txt
io_backend:
type: disk
```
我们使用4个GPU进行训练。还可以使用参数 `--auto_resume` 在必要时自动恢复训练。
```bash
CUDA_VISIBLE_DEVICES=0,1,2,3 \
python -m torch.distributed.launch --nproc_per_node=4 --master_port=4321 realesrgan/train.py -opt options/finetune_realesrgan_x4plus_pairdata.yml --launcher pytorch --auto_resume
```
用 **1个GPU** 训练：
```bash
python realesrgan/train.py -opt options/finetune_realesrgan_x4plus_pairdata.yml --auto_resume
```
================================================
FILE: KDSR-GAN/docs/anime_comparisons.md
================================================
# Comparisons among different anime models
[English](anime_comparisons.md) **|** [简体中文](anime_comparisons_CN.md)
## Update News
- 2022/04/24: Release **AnimeVideo-v3**. We have made the following improvements:
- **better naturalness**
- **Fewer artifacts**
- **more faithful to the original colors**
- **better texture restoration**
- **better background restoration**
## Comparisons
We have compared our RealESRGAN-AnimeVideo-v3 with the following methods.
Our RealESRGAN-AnimeVideo-v3 can achieve better results with faster inference speed.
- [waifu2x](https://github.com/nihui/waifu2x-ncnn-vulkan) with the hyperparameters: `tile=0`, `noiselevel=2`
- [Real-CUGAN](https://github.com/bilibili/ailab/tree/main/Real-CUGAN): we use the [20220227](https://github.com/bilibili/ailab/releases/tag/Real-CUGAN-add-faster-low-memory-mode) version, the hyperparameters are: `cache_mode=0`, `tile=0`, `alpha=1`.
- our RealESRGAN-AnimeVideo-v3
## Results
You may need to **zoom in** for comparing details, or **click the image** to see in the full size. Please note that the images
in the table below are the resized and cropped patches from the original images, you can download the original inputs and outputs from [Google Drive](https://drive.google.com/drive/folders/1bc_Hje1Nqop9NDkUvci2VACSjL7HZMRp?usp=sharing) .
**More natural results, better background restoration**
| Input | waifu2x | Real-CUGAN | RealESRGAN<br>AnimeVideo-v3 |
| :---: | :---: | :---: | :---: |
| |  |  |  |
| |  |  |  |
| |  |  |  |
**Fewer artifacts, better detailed textures**
| Input | waifu2x | Real-CUGAN | RealESRGAN<br>AnimeVideo-v3 |
| :---: | :---: | :---: | :---: |
| |  |  |  |
| |  |  |  |
| |  |  |  |
| |  |  |  |
**Other better results**
| Input | waifu2x | Real-CUGAN | RealESRGAN<br>AnimeVideo-v3 |
| :---: | :---: | :---: | :---: |
| |  |  |  |
| |  |  |  |
|  |   |   |   |
| |  |  |  |
| |  |  |  |
## Inference Speed
### PyTorch
Note that we only report the **model** time, and ignore the IO time.
| GPU | Input Resolution | waifu2x | Real-CUGAN | RealESRGAN-AnimeVideo-v3
| :---: | :---: | :---: | :---: | :---: |
| V100 | 1921 x 1080 | - | 3.4 fps | **10.0** fps |
| V100 | 1280 x 720 | - | 7.2 fps | **22.6** fps |
| V100 | 640 x 480 | - | 24.4 fps | **65.9** fps |
### ncnn
- [ ] TODO
================================================
FILE: KDSR-GAN/docs/anime_comparisons_CN.md
================================================
# 动漫视频模型比较
[English](anime_comparisons.md) **|** [简体中文](anime_comparisons_CN.md)
## 更新
- 2022/04/24: 发布 **AnimeVideo-v3**. 主要做了以下更新：
- **更自然**
- **更少瑕疵**
- **颜色保持得更好**
- **更好的纹理恢复**
- **虚化背景处理**
## 比较
我们将 RealESRGAN-AnimeVideo-v3 与以下方法进行了比较。我们的 RealESRGAN-AnimeVideo-v3 可以以更快的推理速度获得更好的结果。
- [waifu2x](https://github.com/nihui/waifu2x-ncnn-vulkan). 超参数: `tile=0`, `noiselevel=2`
- [Real-CUGAN](https://github.com/bilibili/ailab/tree/main/Real-CUGAN): 我们使用了[20220227](https://github.com/bilibili/ailab/releases/tag/Real-CUGAN-add-faster-low-memory-mode)版本, 超参: `cache_mode=0`, `tile=0`, `alpha=1`.
- 我们的 RealESRGAN-AnimeVideo-v3
## 结果
您可能需要**放大**以比较详细信息, 或者**单击图像**以查看完整尺寸。 请注意下面表格的图片是从原图里裁剪patch并且resize后的结果，您可以从
[Google Drive](https://drive.google.com/drive/folders/1bc_Hje1Nqop9NDkUvci2VACSjL7HZMRp?usp=sharing) 里下载原始的输入和输出。
**更自然的结果，更好的虚化背景恢复**
| 输入 | waifu2x | Real-CUGAN | RealESRGAN<br>AnimeVideo-v3 |
| :---: | :---: | :---: | :---: |
| |  |  |  |
| |  |  |  |
| |  |  |  |
**更少瑕疵，更好的细节纹理**
| 输入 | waifu2x | Real-CUGAN | RealESRGAN<br>AnimeVideo-v3 |
| :---: | :---: | :---: | :---: |
| |  |  |  |
| |  |  |  |
| |  |  |  |
| |  |  |  |
**其他更好的结果**
| 输入 | waifu2x | Real-CUGAN | RealESRGAN<br>AnimeVideo-v3 |
| :---: | :---: | :---: | :---: |
| |  |  |  |
| |  |  |  |
|  |   |   |   |
| |  |  |  |
| |  |  |  |
## 推理速度比较
### PyTorch
请注意，我们只报告了**模型推理**的时间, 而忽略了读写硬盘的时间.
| GPU | 输入尺寸 | waifu2x | Real-CUGAN | RealESRGAN-AnimeVideo-v3
| :---: | :---: | :---: | :---: | :---: |
| V100 | 1921 x 1080 | - | 3.4 fps | **10.0** fps |
| V100 | 1280 x 720 | - | 7.2 fps | **22.6** fps |
| V100 | 640 x 480 | - | 24.4 fps | **65.9** fps |
### ncnn
- [ ] TODO
================================================
FILE: KDSR-GAN/docs/anime_model.md
================================================
# Anime Model
:white_check_mark: We add [*RealESRGAN_x4plus_anime_6B.pth*](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth), which is optimized for **anime** images with much smaller model size.
- [How to Use](#how-to-use)
- [PyTorch Inference](#pytorch-inference)
- [ncnn Executable File](#ncnn-executable-file)
- [Comparisons with waifu2x](#comparisons-with-waifu2x)
- [Comparisons with Sliding Bars](#comparisons-with-sliding-bars)
<p align="center">
<img src="https://raw.githubusercontent.com/xinntao/public-figures/master/Real-ESRGAN/cmp_realesrgan_anime_1.png">
</p>
The following is a video comparison with sliding bar. You may need to use the full-screen mode for better visual quality, as the original image is large; otherwise, you may encounter aliasing issue.
<https://user-images.githubusercontent.com/17445847/131535127-613250d4-f754-4e20-9720-2f9608ad0675.mp4>
## How to Use
### PyTorch Inference
Pre-trained models: [RealESRGAN_x4plus_anime_6B](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth)
```bash
# download model
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth -P experiments/pretrained_models
# inference
python inference_realesrgan.py -n RealESRGAN_x4plus_anime_6B -i inputs
```
### ncnn Executable File
Download the latest portable [Windows](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-windows.zip) / [Linux](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-ubuntu.zip) / [MacOS](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-macos.zip) **executable files for Intel/AMD/Nvidia GPU**.
Taking the Windows as example, run:
```bash
./realesrgan-ncnn-vulkan.exe -i input.jpg -o output.png -n realesrgan-x4plus-anime
```
## Comparisons with waifu2x
We compare Real-ESRGAN-anime with [waifu2x](https://github.com/nihui/waifu2x-ncnn-vulkan). We use the `-n 2 -s 4` for waifu2x.
<p align="center">
<img src="https://raw.githubusercontent.com/xinntao/public-figures/master/Real-ESRGAN/cmp_realesrgan_anime_1.png">
</p>
<p align="center">
<img src="https://raw.githubusercontent.com/xinntao/public-figures/master/Real-ESRGAN/cmp_realesrgan_anime_2.png">
</p>
<p align="center">
<img src="https://raw.githubusercontent.com/xinntao/public-figures/master/Real-ESRGAN/cmp_realesrgan_anime_3.png">
</p>
<p align="center">
<img src="https://raw.githubusercontent.com/xinntao/public-figures/master/Real-ESRGAN/cmp_realesrgan_anime_4.png">
</p>
<p align="center">
<img src="https://raw.githubusercontent.com/xinntao/public-figures/master/Real-ESRGAN/cmp_realesrgan_anime_5.png">
</p>
## Comparisons with Sliding Bars
The following are video comparisons with sliding bar. You may need to use the full-screen mode for better visual quality, as the original image is large; otherwise, you may encounter aliasing issue.
<https://user-images.githubusercontent.com/17445847/131536647-a2fbf896-b495-4a9f-b1dd-ca7bbc90101a.mp4>
<https://user-images.githubusercontent.com/17445847/131536742-6d9d82b6-9765-4296-a15f-18f9aeaa5465.mp4>
================================================
FILE: KDSR-GAN/docs/anime_video_model.md
================================================
# Anime Video Models
:white_check_mark: We add small models that are optimized for anime videos :-)<br>
More comparisons can be found in [anime_comparisons.md](anime_comparisons.md)
- [How to Use](#how-to-use)
- [PyTorch Inference](#pytorch-inference)
- [ncnn Executable File](#ncnn-executable-file)
- [Step 1: Use ffmpeg to extract frames from video](#step-1-use-ffmpeg-to-extract-frames-from-video)
- [Step 2: Inference with Real-ESRGAN executable file](#step-2-inference-with-real-esrgan-executable-file)
- [Step 3: Merge the enhanced frames back into a video](#step-3-merge-the-enhanced-frames-back-into-a-video)
- [More Demos](#more-demos)
| Models | Scale | Description |
| ---------------------------------------------------------------------------------------------------------------------------------- | :---- | :----------------------------- |
| [realesr-animevideov3](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth) | X4 <sup>1</sup> | Anime video model with XS size |
Note: <br>
<sup>1</sup> This model can also be used for X1, X2, X3.
---
The following are some demos (best view in the full screen mode).
<https://user-images.githubusercontent.com/17445847/145706977-98bc64a4-af27-481c-8abe-c475e15db7ff.MP4>
<https://user-images.githubusercontent.com/17445847/145707055-6a4b79cb-3d9d-477f-8610-c6be43797133.MP4>
<https://user-images.githubusercontent.com/17445847/145783523-f4553729-9f03-44a8-a7cc-782aadf67b50.MP4>
## How to Use
### PyTorch Inference
```bash
# download model
wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth -P realesrgan/weights
# single gpu and single process inference
CUDA_VISIBLE_DEVICES=0 python inference_realesrgan_video.py -i inputs/video/onepiece_demo.mp4 -n realesr-animevideov3 -s 2 --suffix outx2
# single gpu and multi process inference (you can use multi-processing to improve GPU utilization)
CUDA_VISIBLE_DEVICES=0 python inference_realesrgan_video.py -i inputs/video/onepiece_demo.mp4 -n realesr-animevideov3 -s 2 --suffix outx2 --num_process_per_gpu 2
# multi gpu and multi process inference
CUDA_VISIBLE_DEVICES=0,1,2,3 python inference_realesrgan_video.py -i inputs/video/onepiece_demo.mp4 -n realesr-animevideov3 -s 2 --suffix outx2 --num_process_per_gpu 2
```
```console
Usage:
--num_process_per_gpu The total number of process is num_gpu * num_process_per_gpu. The bottleneck of
the program lies on the IO, so the GPUs are usually not fully utilized. To alleviate
this issue, you can use multi-processing by setting this parameter. As long as it
does not exceed the CUDA memory
--extract_frame_first If you encounter ffmpeg error when using multi-processing, you can turn this option on.
```
### NCNN Executable File
#### Step 1: Use ffmpeg to extract frames from video
```bash
ffmpeg -i onepiece_demo.mp4 -qscale:v 1 -qmin 1 -qmax 1 -vsync 0 tmp_frames/frame%08d.png
```
- Remember to create the folder `tmp_frames` ahead
#### Step 2: Inference with Real-ESRGAN executable file
1. Download the latest portable [Windows](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-windows.zip) / [Linux](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-ubuntu.zip) / [MacOS](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-macos.zip) **executable files for Intel/AMD/Nvidia GPU**
1. Taking the Windows as example, run:
```bash
./realesrgan-ncnn-vulkan.exe -i tmp_frames -o out_frames -n realesr-animevideov3 -s 2 -f jpg
```
- Remember to create the folder `out_frames` ahead
#### Step 3: Merge the enhanced frames back into a video
1. First obtain fps from input videos by
```bash
ffmpeg -i onepiece_demo.mp4
```
```console
Usage:
-i input video path
```
You will get the output similar to the following screenshot.
<p align="center">
<img src="https://user-images.githubusercontent.com/17445847/145710145-c4f3accf-b82f-4307-9f20-3803a2c73f57.png">
</p>
2. Merge frames
```bash
ffmpeg -r 23.98 -i out_frames/frame%08d.jpg -c:v libx264 -r 23.98 -pix_fmt yuv420p output.mp4
```
```console
Usage:
-i input video path
-c:v video encoder (usually we use libx264)
-r fps, remember to modify it to meet your needs
-pix_fmt pixel format in video
```
If you also want to copy audio from the input videos, run:
```bash
ffmpeg -r 23.98 -i out_frames/frame%08d.jpg -i onepiece_demo.mp4 -map 0:v:0 -map 1:a:0 -c:a copy -c:v libx264 -r 23.98 -pix_fmt yuv420p output_w_audio.mp4
```
```console
Usage:
-i input video path, here we use two input streams
-c:v video encoder (usually we use libx264)
-r fps, remember to modify it to meet your needs
-pix_fmt pixel format in video
```
## More Demos
- Input video for One Piece:
<https://user-images.githubusercontent.com/17445847/145706822-0e83d9c4-78ef-40ee-b2a4-d8b8c3692d17.mp4>
- Out video for One Piece
<https://user-images.githubusercontent.com/17445847/164960481-759658cf-fcb8-480c-b888-cecb606e8744.mp4>
**More comparisons**
<https://user-images.githubusercontent.com/17445847/145707458-04a5e9b9-2edd-4d1f-b400-380a72e5f5e6.MP4>
================================================
FILE: KDSR-GAN/docs/feedback.md
================================================
# Feedback 反馈
## 动漫插画模型
1. 视频处理不了: 目前的模型，不是针对视频的，所以视频效果很很不好。我们在探究针对视频的模型了
1. 景深虚化有问题: 现在的模型把一些景深 和 特意的虚化 都复原了，感觉不好。这个后面我们会考虑把这个信息结合进入。一个简单的做法是识别景深和虚化，然后作为条件告诉神经网络，哪些地方复原强一些，哪些地方复原要弱一些
1. 不可以调节: 像 Waifu2X 可以调节。可以根据自己的喜好，做调整，但是 Real-ESRGAN-anime 并不可以。导致有些恢复效果过了
1. 把原来的风格改变了: 不同的动漫插画都有自己的风格，现在的 Real-ESRGAN-anime 倾向于恢复成一种风格（这是受到训练数据集影响的）。风格是动漫很重要的一个要素，所以要尽可能保持
1. 模型太大: 目前的模型处理太慢，能够更快。这个我们有相关的工作在探究，希望能够尽快有结果，并应用到 Real-ESRGAN 这一系列的模型上
Thanks for the [detailed and valuable feedbacks/suggestions](https://github.com/xinntao/Real-ESRGAN/issues/131) by [2ji3150](https://github.com/2ji3150).
================================================
FILE: KDSR-GAN/docs/model_zoo.md
================================================
# :european_castle: Model Zoo
- [For General Images](#for-general-images)
- [For Anime Images](#for-anime-images)
- [For Anime Videos](#for-anime-videos)
---
## For General Images
| Models | Scale | Description |
| ------------------------------------------------------------------------------------------------------------------------------- | :---- | :------------------------------------------- |
| [RealESRGAN_x4plus](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth) | X4 | X4 model for general images |
| [RealESRGAN_x2plus](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth) | X2 | X2 model for general images |
| [RealESRNet_x4plus](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/RealESRNet_x4plus.pth) | X4 | X4 model with MSE loss (over-smooth effects) |
| [official ESRGAN_x4](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.1/ESRGAN_SRx4_DF2KOST_official-ff704c30.pth) | X4 | official ESRGAN model |
The following models are **discriminators**, which are usually used for fine-tuning.
| Models | Corresponding model |
| ---------------------------------------------------------------------------------------------------------------------- | :------------------ |
| [RealESRGAN_x4plus_netD](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.3/RealESRGAN_x4plus_netD.pth) | RealESRGAN_x4plus |
| [RealESRGAN_x2plus_netD](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.3/RealESRGAN_x2plus_netD.pth) | RealESRGAN_x2plus |
## For Anime Images / Illustrations
| Models | Scale | Description |
| ------------------------------------------------------------------------------------------------------------------------------ | :---- | :---------------------------------------------------------- |
| [RealESRGAN_x4plus_anime_6B](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth) | X4 | Optimized for anime images; 6 RRDB blocks (smaller network) |
The following models are **discriminators**, which are usually used for fine-tuning.
| Models | Corresponding model |
| ---------------------------------------------------------------------------------------------------------------------------------------- | :------------------------- |
| [RealESRGAN_x4plus_anime_6B_netD](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B_netD.pth) | RealESRGAN_x4plus_anime_6B |
## For Animation Videos
| Models | Scale | Description |
| ---------------------------------------------------------------------------------------------------------------------------------- | :---- | :----------------------------- |
| [realesr-animevideov3](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-animevideov3.pth) | X4<sup>1</sup> | Anime video model with XS size |
Note: <br>
<sup>1</sup> This model can also be used for X1, X2, X3.
The following models are **discriminators**, which are usually used for fine-tuning.
TODO
================================================
FILE: KDSR-GAN/docs/ncnn_conversion.md
================================================
# Instructions on converting to NCNN models
1. Convert to onnx model with `scripts/pytorch2onnx.py`. Remember to modify codes accordingly
1. Convert onnx model to ncnn model
1. `cd ncnn-master\ncnn\build\tools\onnx`
1. `onnx2ncnn.exe realesrgan-x4.onnx realesrgan-x4-raw.param realesrgan-x4-raw.bin`
1. Optimize ncnn model
1. fp16 mode
1. `cd ncnn-master\ncnn\build\tools`
1. `ncnnoptimize.exe realesrgan-x4-raw.param realesrgan-x4-raw.bin realesrgan-x4.param realesrgan-x4.bin 1`
1. Modify the blob name in `realesrgan-x4.param`: `data` and `output`
================================================
FILE: KDSR-GAN/kdsrgan/__init__.py
================================================
# flake8: noqa
from .losses import *
from .archs import *
from .data import *
from .models import *
from .utils import *
from .version import *
================================================
FILE: KDSR-GAN/kdsrgan/archs/ST_arch.py
================================================
from basicsr.utils.registry import ARCH_REGISTRY
from torch.nn import functional as F
import torch
from torch import nn
import kdsrgan.archs.common as common
class IDR_DDC(nn.Module):
def __init__(self, channels_in, channels_out, kernel_size, reduction):
super(IDR_DDC, self).__init__()
self.channels_out = channels_out
self.channels_in = channels_in
self.kernel_size = kernel_size
self.kernel = nn.Sequential(
nn.Linear(channels_in, channels_in, bias=False),
nn.LeakyReLU(0.1, True),
nn.Linear(channels_in, channels_in * self.kernel_size * self.kernel_size, bias=False)
)
def forward(self, x):
'''
:param x[0]: feature map: B * C * H * W
:param x[1]: degradation representation: B * C
'''
b, c, h, w = x[0].size()
# branch 1
kernel = self.kernel(x[1]).view(-1, 1, self.kernel_size, self.kernel_size)
out = F.conv2d(x[0].view(1, -1, h, w), kernel, groups=b*c, padding=(self.kernel_size-1)//2)
out = out.view(b, -1, h, w)
return out
class IDR_DCRB(nn.Module):
def __init__(self, conv, n_feat, kernel_size, reduction):
super(IDR_DCRB, self).__init__()
self.da_conv1 = IDR_DDC(n_feat, n_feat, kernel_size, reduction)
self.conv1 = conv(n_feat, n_feat, kernel_size)
self.relu = nn.LeakyReLU(0.1, True)
def forward(self, x):
'''
:param x[0]: feature map: B * C * H * W
:param x[1]: degradation representation: B * C
'''
out = self.relu(self.da_conv1(x))
out = self.conv1(out)
out = out + x[0]
return out
class DAG(nn.Module):
def __init__(self, conv, n_feat, kernel_size, reduction, n_blocks):
super(DAG, self).__init__()
self.n_blocks = n_blocks
modules_body = [
IDR_DCRB(conv, n_feat, kernel_size, reduction) \
for _ in range(n_blocks)
]
# modules_body.append(conv(n_feat, n_feat, kernel_size))
self.body = nn.Sequential(*modules_body)
def forward(self, x):
'''
:param x[0]: feature map: B * C * H * W
:param x[1]: degradation representation: B * C
'''
res = x[0]
for i in range(self.n_blocks):
res = self.body[i]([res, x[1]])
return res
class KDSR(nn.Module):
def __init__(self, n_feats=128, scale=4,n_sr_blocks = 42, rgb_range=1,reduction = 8, conv=common.default_conv):
super(KDSR, self).__init__()
kernel_size = 3
# RGB mean for DIV2K
rgb_mean = (0.4488, 0.4371, 0.4040)
rgb_std = (1.0, 1.0, 1.0)
self.sub_mean = common.MeanShift(rgb_range, rgb_mean, rgb_std)
self.add_mean = common.MeanShift(rgb_range, rgb_mean, rgb_std, 1)
# head module
modules_head = [conv(3, n_feats, kernel_size)]
self.head = nn.Sequential(*modules_head)
# body
modules_body = [
DAG(common.default_conv, n_feats, kernel_size, reduction, n_sr_blocks)
]
modules_body.append(conv(n_feats, n_feats, kernel_size))
self.body = nn.Sequential(*modules_body)
# tail
modules_tail = [common.Upsampler(conv, scale, n_feats, act=False),
conv(n_feats, 3, kernel_size)]
self.tail = nn.Sequential(*modules_tail)
def forward(self, x, k_v):
# sub mean
x = self.sub_mean(x)
# head
x = self.head(x)
# body
res = x
res = self.body[0]([res, k_v])
res = self.body[-1](res)
res = res + x
# tail
x = self.tail(res)
# add mean
x = self.add_mean(x)
return x
class KD_IDE(nn.Module):
def __init__(self,n_feats = 128, n_encoder_res = 6):
super(KD_IDE, self).__init__()
E1=[nn.Conv2d(3, n_feats, kernel_size=3, padding=1),
nn.LeakyReLU(0.1, True)]
E2=[
common.ResBlock(
common.default_conv, n_feats, kernel_size=3
) for _ in range(n_encoder_res)
]
E3=[
nn.Conv2d(n_feats, n_feats * 2, kernel_size=3, padding=1),
nn.LeakyReLU(0.1, True),
nn.Conv2d(n_feats * 2, n_feats * 2, kernel_size=3, padding=1),
nn.LeakyReLU(0.1, True),
nn.Conv2d(n_feats * 2, n_feats * 4, kernel_size=3, padding=1),
nn.LeakyReLU(0.1, True),
nn.AdaptiveAvgPool2d(1),
]
E=E1+E2+E3
self.E = nn.Sequential(
*E
)
self.mlp = nn.Sequential(
nn.Linear(n_feats * 4, n_feats * 4),
nn.LeakyReLU(0.1, True),
nn.Linear(n_feats * 4, n_feats * 4),
nn.LeakyReLU(0.1, True)
)
self.compress = nn.Sequential(
nn.Linear(n_feats*4, n_feats),
nn.LeakyReLU(0.1, True)
)
def forward(self, x):
fea = self.E(x).squeeze(-1).squeeze(-1)
S_fea = []
fea1 = self.mlp(fea)
fea = self.compress(fea1)
S_fea.append(fea1)
return fea,S_fea
@ARCH_REGISTRY.register()
class BlindSR_ST(nn.Module):
def __init__(self, n_feats=128, n_encoder_res=6, scale=4,n_sr_blocks=42 ):
super(BlindSR_ST, self).__init__()
# Generator
self.G = KDSR(n_feats=n_feats, scale=scale,n_sr_blocks=n_sr_blocks)
self.E_st = KD_IDE(n_feats=n_feats, n_encoder_res=n_encoder_res)
self.pixel_unshuffle = nn.PixelUnshuffle(scale)
def forward(self, x):
if self.training:
# degradation-aware represenetion learning
deg_repre, S_fea = self.E_st(x)
# degradation-aware SR
sr = self.G(x, deg_repre)
return sr, S_fea
else:
# degradation-aware represenetion learning
deg_repre, _ = self.E_st(x)
# degradation-aware SR
sr = self.G(x, deg_repre)
return sr
================================================
FILE: KDSR-GAN/kdsrgan/archs/TA_arch.py
================================================
from basicsr.utils.registry import ARCH_REGISTRY
from torch.nn import functional as F
import torch
from torch import nn
import kdsrgan.archs.common as common
class IDR_DDC(nn.Module):
def __init__(self, channels_in, channels_out, kernel_size, reduction):
super(IDR_DDC, self).__init__()
self.channels_out = channels_out
self.channels_in = channels_in
self.kernel_size = kernel_size
self.kernel = nn.Sequential(
nn.Linear(channels_in, channels_in, bias=False),
nn.LeakyReLU(0.1, True),
nn.Linear(channels_in, channels_in * self.kernel_size * self.kernel_size, bias=False)
)
def forward(self, x):
'''
:param x[0]: feature map: B * C * H * W
:param x[1]: degradation representation: B * C
'''
b, c, h, w = x[0].size()
# branch 1
kernel = self.kernel(x[1]).view(-1, 1, self.kernel_size, self.kernel_size)
out = F.conv2d(x[0].view(1, -1, h, w), kernel, groups=b*c, padding=(self.kernel_size-1)//2)
out = out.view(b, -1, h, w)
return out
class IDR_DCRB(nn.Module):
def __init__(self, conv, n_feat, kernel_size, reduction):
super(IDR_DCRB, self).__init__()
self.da_conv1 = IDR_DDC(n_feat, n_feat, kernel_size, reduction)
self.conv1 = conv(n_feat, n_feat, kernel_size)
self.relu = nn.LeakyReLU(0.1, True)
def forward(self, x):
'''
:param x[0]: feature map: B * C * H * W
:param x[1]: degradation representation: B * C
'''
out = self.relu(self.da_conv1(x))
out = self.conv1(out)
out = out + x[0]
return out
class DAG(nn.Module):
def __init__(self, conv, n_feat, kernel_size, reduction, n_blocks):
super(DAG, self).__init__()
self.n_blocks = n_blocks
modules_body = [
IDR_DCRB(conv, n_feat, kernel_size, reduction) \
for _ in range(n_blocks)
]
# modules_body.append(conv(n_feat, n_feat, kernel_size))
self.body = nn.Sequential(*modules_body)
def forward(self, x):
'''
:param x[0]: feature map: B * C * H * W
:param x[1]: degradation representation: B * C
'''
res = x[0]
for i in range(self.n_blocks):
res = self.body[i]([res, x[1]])
# res = self.body[-1](res)
#res = res + x[0]
return res
class KDSR(nn.Module):
def __init__(self, n_feats=128, scale=4,n_sr_blocks = 42, rgb_range=1,reduction = 8, conv=common.default_conv):
super(KDSR, self).__init__()
kernel_size = 3
# RGB mean for DIV2K
rgb_mean = (0.4488, 0.4371, 0.4040)
rgb_std = (1.0, 1.0, 1.0)
self.sub_mean = common.MeanShift(rgb_range, rgb_mean, rgb_std)
self.add_mean = common.MeanShift(rgb_range, rgb_mean, rgb_std, 1)
# head module
modules_head = [conv(3, n_feats, kernel_size)]
self.head = nn.Sequential(*modules_head)
# body
modules_body = [
DAG(common.default_conv, n_feats, kernel_size, reduction, n_sr_blocks)
]
modules_body.append(conv(n_feats, n_feats, kernel_size))
self.body = nn.Sequential(*modules_body)
# tail
modules_tail = [common.Upsampler(conv, scale, n_feats, act=False),
conv(n_feats, 3, kernel_size)]
self.tail = nn.Sequential(*modules_tail)
def forward(self, x, k_v):
# sub mean
x = self.sub_mean(x)
# head
x = self.head(x)
# body
res = x
res = self.body[0]([res, k_v])
res = self.body[-1](res)
res = res + x
# tail
x = self.tail(res)
# add mean
x = self.add_mean(x)
return x
class KD_IDE(nn.Module):
def __init__(self,n_feats = 128, n_encoder_res = 6):
super(KD_IDE, self).__init__()
E1=[nn.Conv2d(51, n_feats, kernel_size=3, padding=1),
nn.LeakyReLU(0.1, True)]
E2=[
common.ResBlock(
common.default_conv, n_feats, kernel_size=3
) for _ in range(n_encoder_res)
]
E3=[
nn.Conv2d(n_feats, n_feats * 2, kernel_size=3, padding=1),
nn.LeakyReLU(0.1, True),
nn.Conv2d(n_feats * 2, n_feats * 2, kernel_size=3, padding=1),
nn.LeakyReLU(0.1, True),
nn.Conv2d(n_feats * 2, n_feats * 4, kernel_size=3, padding=1),
nn.LeakyReLU(0.1, True),
nn.AdaptiveAvgPool2d(1),
]
E=E1+E2+E3
self.E = nn.Sequential(
*E
)
self.mlp = nn.Sequential(
nn.Linear(n_feats * 4, n_feats * 4),
nn.LeakyReLU(0.1, True),
nn.Linear(n_feats * 4, n_feats * 4),
nn.LeakyReLU(0.1, True)
)
self.compress = nn.Sequential(
nn.Linear(n_feats*4, n_feats),
nn.LeakyReLU(0.1, True)
)
def forward(self, x):
fea = self.E(x).squeeze(-1).squeeze(-1)
T_fea = []
fea1 = self.mlp(fea)
fea = self.compress(fea1)
T_fea.append(fea1)
return fea,T_fea
@ARCH_REGISTRY.register()
class BlindSR_TA(nn.Module):
def __init__(self, n_feats=128, n_encoder_res=6, scale=4,n_sr_blocks=42 ):
super(BlindSR_TA, self).__init__()
# Generator
self.G = KDSR(n_feats=n_feats, scale=scale,n_sr_blocks=n_sr_blocks)
self.E = KD_IDE(n_feats=n_feats, n_encoder_res=n_encoder_res)
self.pixel_unshuffle = nn.PixelUnshuffle(scale)
def forward(self, x, gt):
if self.training:
hr = self.pixel_unshuffle(gt)
deg_repre = torch.cat([x, hr], dim=1)
# degradation-aware represenetion learning
deg_repre, T_fea = self.E(deg_repre)
# degradation-aware SR
sr = self.G(x, deg_repre)
return sr, T_fea
else:
# degradation-aware represenetion learning
hr = self.pixel_unshuffle(gt)
deg_repre = torch.cat([x, hr], dim=1)
deg_repre, _ = self.E(deg_repre)
# degradation-aware SR
sr = self.G(x, deg_repre)
return sr
================================================
FILE: KDSR-GAN/kdsrgan/archs/__init__.py
================================================
import importlib
from basicsr.utils import scandir
from os import path as osp
# automatically scan and import arch modules for registry
# scan all the files that end with '_arch.py' under the archs folder
arch_folder = osp.dirname(osp.abspath(__file__))
arch_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(arch_folder) if v.endswith('_arch.py')]
# import all the arch modules
_arch_modules = [importlib.import_module(f'kdsrgan.archs.{file_name}') for file_name in arch_filenames]
================================================
FILE: KDSR-GAN/kdsrgan/archs/common.py
================================================
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
def default_conv(in_channels, out_channels, kernel_size, bias=True):
return nn.Conv2d(in_channels, out_channels, kernel_size, padding=(kernel_size//2), bias=bias)
class ResBlock(nn.Module):
def __init__(
self, conv, n_feats, kernel_size,
bias=True, bn=False, act=nn.LeakyReLU(0.1, inplace=True), res_scale=1):
super(ResBlock, self).__init__()
m = []
for i in range(2):
m.append(conv(n_feats, n_feats, kernel_size, bias=bias))
if bn:
m.append(nn.BatchNorm2d(n_feats))
if i == 0:
m.append(act)
self.body = nn.Sequential(*m)
# self.res_scale = res_scale
def forward(self, x):
res = self.body(x)
res += x
return res
class MeanShift(nn.Conv2d):
def __init__(self, rgb_range, rgb_mean, rgb_std, sign=-1):
super(MeanShift, self).__init__(3, 3, kernel_size=1)
std = torch.Tensor(rgb_std)
self.weight.data = torch.eye(3).view(3, 3, 1, 1)
self.weight.data.div_(std.view(3, 1, 1, 1))
self.bias.data = sign * rgb_range * torch.Tensor(rgb_mean)
self.bias.data.div_(std)
self.weight.requires_grad = False
self.bias.requires_grad = False
class Upsampler(nn.Sequential):
def __init__(self, conv, scale, n_feat, act=False, bias=True):
m = []
if (int(scale) & (int(scale) - 1)) == 0: # Is scale = 2^n?
for _ in range(int(math.log(scale, 2))):
m.append(conv(n_feat, 4 * n_feat, 3, bias))
m.append(nn.PixelShuffle(2))
if act: m.append(act())
elif scale == 3:
m.append(conv(n_feat, 9 * n_feat, 3, bias))
m.append(nn.PixelShuffle(3))
if act: m.append(act())
else:
raise NotImplementedError
super(Upsampler, self).__init__(*m)
================================================
FILE: KDSR-GAN/kdsrgan/archs/discriminator_arch.py
================================================
from basicsr.utils.registry import ARCH_REGISTRY
from torch import nn as nn
from torch.nn import functional as F
from torch.nn.utils import spectral_norm
@ARCH_REGISTRY.register()
class UNetDiscriminatorSN(nn.Module):
"""Defines a U-Net discriminator with spectral normalization (SN)
It is used in Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
Arg:
num_in_ch (int): Channel number of inputs. Default: 3.
num_feat (int): Channel number of base intermediate features. Default: 64.
skip_connection (bool): Whether to use skip connections between U-Net. Default: True.
"""
def __init__(self, num_in_ch, num_feat=64, skip_connection=True):
super(UNetDiscriminatorSN, self).__init__()
self.skip_connection = skip_connection
norm = spectral_norm
# the first convolution
self.conv0 = nn.Conv2d(num_in_ch, num_feat, kernel_size=3, stride=1, padding=1)
# downsample
self.conv1 = norm(nn.Conv2d(num_feat, num_feat * 2, 4, 2, 1, bias=False))
self.conv2 = norm(nn.Conv2d(num_feat * 2, num_feat * 4, 4, 2, 1, bias=False))
self.conv3 = norm(nn.Conv2d(num_feat * 4, num_feat * 8, 4, 2, 1, bias=False))
# upsample
self.conv4 = norm(nn.Conv2d(num_feat * 8, num_feat * 4, 3, 1, 1, bias=False))
self.conv5 = norm(nn.Conv2d(num_feat * 4, num_feat * 2, 3, 1, 1, bias=False))
self.conv6 = norm(nn.Conv2d(num_feat * 2, num_feat, 3, 1, 1, bias=False))
# extra convolutions
self.conv7 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))
self.conv8 = norm(nn.Conv2d(num_feat, num_feat, 3, 1, 1, bias=False))
self.conv9 = nn.Conv2d(num_feat, 1, 3, 1, 1)
def forward(self, x):
# downsample
x0 = F.leaky_relu(self.conv0(x), negative_slope=0.2, inplace=True)
x1 = F.leaky_relu(self.conv1(x0), negative_slope=0.2, inplace=True)
x2 = F.leaky_relu(self.conv2(x1), negative_slope=0.2, inplace=True)
x3 = F.leaky_relu(self.conv3(x2), negative_slope=0.2, inplace=True)
# upsample
x3 = F.interpolate(x3, scale_factor=2, mode='bilinear', align_corners=False)
x4 = F.leaky_relu(self.conv4(x3), negative_slope=0.2, inplace=True)
if self.skip_connection:
x4 = x4 + x2
x4 = F.interpolate(x4, scale_factor=2, mode='bilinear', align_corners=False)
x5 = F.leaky_relu(self.conv5(x4), negative_slope=0.2, inplace=True)
if self.skip_connection:
x5 = x5 + x1
x5 = F.interpolate(x5, scale_factor=2, mode='bilinear', align_corners=False)
x6 = F.leaky_relu(self.conv6(x5), negative_slope=0.2, inplace=True)
if self.skip_connection:
x6 = x6 + x0
# extra convolutions
out = F.leaky_relu(self.conv7(x6), negative_slope=0.2, inplace=True)
out = F.leaky_relu(self.conv8(out), negative_slope=0.2, inplace=True)
out = self.conv9(out)
return out
================================================
FILE: KDSR-GAN/kdsrgan/archs/srvgg_arch.py
================================================
from basicsr.utils.registry import ARCH_REGISTRY
from torch import nn as nn
from torch.nn import functional as F
@ARCH_REGISTRY.register()
class SRVGGNetCompact(nn.Module):
"""A compact VGG-style network structure for super-resolution.
It is a compact network structure, which performs upsampling in the last layer and no convolution is
conducted on the HR feature space.
Args:
num_in_ch (int): Channel number of inputs. Default: 3.
num_out_ch (int): Channel number of outputs. Default: 3.
num_feat (int): Channel number of intermediate features. Default: 64.
num_conv (int): Number of convolution layers in the body network. Default: 16.
upscale (int): Upsampling factor. Default: 4.
act_type (str): Activation type, options: 'relu', 'prelu', 'leakyrelu'. Default: prelu.
"""
def __init__(self, num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16, upscale=4, act_type='prelu'):
super(SRVGGNetCompact, self).__init__()
self.num_in_ch = num_in_ch
self.num_out_ch = num_out_ch
self.num_feat = num_feat
self.num_conv = num_conv
self.upscale = upscale
self.act_type = act_type
self.body = nn.ModuleList()
# the first conv
self.body.append(nn.Conv2d(num_in_ch, num_feat, 3, 1, 1))
# the first activation
if act_type == 'relu':
activation = nn.ReLU(inplace=True)
elif act_type == 'prelu':
activation = nn.PReLU(num_parameters=num_feat)
elif act_type == 'leakyrelu':
activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)
self.body.append(activation)
# the body structure
for _ in range(num_conv):
self.body.append(nn.Conv2d(num_feat, num_feat, 3, 1, 1))
# activation
if act_type == 'relu':
activation = nn.ReLU(inplace=True)
elif act_type == 'prelu':
activation = nn.PReLU(num_parameters=num_feat)
elif act_type == 'leakyrelu':
activation = nn.LeakyReLU(negative_slope=0.1, inplace=True)
self.body.append(activation)
# the last conv
self.body.append(nn.Conv2d(num_feat, num_out_ch * upscale * upscale, 3, 1, 1))
# upsample
self.upsampler = nn.PixelShuffle(upscale)
def forward(self, x):
out = x
for i in range(0, len(self.body)):
out = self.body[i](out)
out = self.upsampler(out)
# add the nearest upsampled image, so that the network learns the residual
base = F.interpolate(x, scale_factor=self.upscale, mode='nearest')
out += base
return out
================================================
FILE: KDSR-GAN/kdsrgan/data/__init__.py
================================================
import importlib
from basicsr.utils import scandir
from os import path as osp
# automatically scan and import dataset modules for registry
# scan all the files that end with '_dataset.py' under the data folder
data_folder = osp.dirname(osp.abspath(__file__))
dataset_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(data_folder) if v.endswith('_dataset.py')]
# import all the dataset modules
_dataset_modules = [importlib.import_module(f'kdsrgan.data.{file_name}') for file_name in dataset_filenames]
================================================
FILE: KDSR-GAN/kdsrgan/data/kdsrgan_dataset.py
================================================
import cv2
import math
import numpy as np
import os
import os.path as osp
import random
import time
import torch
from basicsr.data.degradations import circular_lowpass_kernel, random_mixed_kernels
from basicsr.data.transforms import augment
from basicsr.utils import FileClient, get_root_logger, imfrombytes, img2tensor
from basicsr.utils.registry import DATASET_REGISTRY
from torch.utils import data as data
@DATASET_REGISTRY.register()
class KDSRGANDataset(data.Dataset):
"""Dataset used for KDSRGAN model:
KDSRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data.
It loads gt (Ground-Truth) images, and augments them.
It also generates blur kernels and sinc kernels for generating low-quality images.
Note that the low-quality images are processed in tensors on GPUS for faster processing.
Args:
opt (dict): Config for train datasets. It contains the following keys:
dataroot_gt (str): Data root path for gt.
meta_info (str): Path for meta information file.
io_backend (dict): IO backend type and other kwarg.
use_hflip (bool): Use horizontal flips.
use_rot (bool): Use rotation (use vertical flip and transposing h and w for implementation).
Please see more options in the codes.
"""
def __init__(self, opt):
super(KDSRGANDataset, self).__init__()
self.opt = opt
self.file_client = None
self.io_backend_opt = opt['io_backend']
self.gt_folder = opt['dataroot_gt']
# file client (lmdb io backend)
if self.io_backend_opt['type'] == 'lmdb':
self.io_backend_opt['db_paths'] = [self.gt_folder]
self.io_backend_opt['client_keys'] = ['gt']
if not self.gt_folder.endswith('.lmdb'):
raise ValueError(f"'dataroot_gt' should end with '.lmdb', but received {self.gt_folder}")
with open(osp.join(self.gt_folder, 'meta_info.txt')) as fin:
self.paths = [line.split('.')[0] for line in fin]
else:
# disk backend with meta_info
# Each line in the meta_info describes the relative path to an image
with open(self.opt['meta_info']) as fin:
paths = [line.strip().split(' ')[0] for line in fin]
self.paths = [os.path.join(self.gt_folder, v) for v in paths]
# blur settings for the first degradation
self.blur_kernel_size = opt['blur_kernel_size']
self.kernel_list = opt['kernel_list']
self.kernel_prob = opt['kernel_prob'] # a list for each kernel probability
self.blur_sigma = opt['blur_sigma']
self.betag_range = opt['betag_range'] # betag used in generalized Gaussian blur kernels
self.betap_range = opt['betap_range'] # betap used in plateau blur kernels
self.sinc_prob = opt['sinc_prob'] # the probability for sinc filters
# blur settings for the second degradation
self.blur_kernel_size2 = opt['blur_kernel_size2']
self.kernel_list2 = opt['kernel_list2']
self.kernel_prob2 = opt['kernel_prob2']
self.blur_sigma2 = opt['blur_sigma2']
self.betag_range2 = opt['betag_range2']
self.betap_range2 = opt['betap_range2']
self.sinc_prob2 = opt['sinc_prob2']
# a final sinc filter
self.final_sinc_prob = opt['final_sinc_prob']
self.kernel_range = [2 * v + 1 for v in range(3, 11)] # kernel size ranges from 7 to 21
# TODO: kernel range is now hard-coded, should be in the configure file
self.pulse_tensor = torch.zeros(21, 21).float() # convolving with pulse tensor brings no blurry effect
self.pulse_tensor[10, 10] = 1
def __getitem__(self, index):
if self.file_client is None:
self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
# -------------------------------- Load gt images -------------------------------- #
# Shape: (h, w, c); channel order: BGR; image range: [0, 1], float32.
gt_path = self.paths[index]
# avoid errors caused by high latency in reading files
retry = 3
while retry > 0:
try:
img_bytes = self.file_client.get(gt_path, 'gt')
except (IOError, OSError) as e:
logger = get_root_logger()
logger.warn(f'File client error: {e}, remaining retry times: {retry - 1}')
# change another file to read
index = random.randint(0, self.__len__())
gt_path = self.paths[index]
time.sleep(1) # sleep 1s for occasional server congestion
else:
break
finally:
retry -= 1
img_gt = imfrombytes(img_bytes, float32=True)
# -------------------- Do augmentation for training: flip, rotation -------------------- #
img_gt = augment(img_gt, self.opt['use_hflip'], self.opt['use_rot'])
# crop or pad to 400
# TODO: 400 is hard-coded. You may change it accordingly
h, w = img_gt.shape[0:2]
crop_pad_size = 400
# pad
if h < crop_pad_size or w < crop_pad_size:
pad_h = max(0, crop_pad_size - h)
pad_w = max(0, crop_pad_size - w)
img_gt = cv2.copyMakeBorder(img_gt, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT_101)
# crop
if img_gt.shape[0] > crop_pad_size or img_gt.shape[1] > crop_pad_size:
h, w = img_gt.shape[0:2]
# randomly choose top and left coordinates
top = random.randint(0, h - crop_pad_size)
left = random.randint(0, w - crop_pad_size)
img_gt = img_gt[top:top + crop_pad_size, left:left + crop_pad_size, ...]
# ------------------------ Generate kernels (used in the first degradation) ------------------------ #
kernel_size = random.choice(self.kernel_range)
if np.random.uniform() < self.opt['sinc_prob']:
# this sinc filter setting is for kernels ranging from [7, 21]
if kernel_size < 13:
omega_c = np.random.uniform(np.pi / 3, np.pi)
else:
omega_c = np.random.uniform(np.pi / 5, np.pi)
kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False)
else:
kernel = random_mixed_kernels(
self.kernel_list,
self.kernel_prob,
kernel_size,
self.blur_sigma,
self.blur_sigma, [-math.pi, math.pi],
self.betag_range,
self.betap_range,
noise_range=None)
# pad kernel
pad_size = (21 - kernel_size) // 2
kernel = np.pad(kernel, ((pad_size, pad_size), (pad_size, pad_size)))
# ------------------------ Generate kernels (used in the second degradation) ------------------------ #
kernel_size = random.choice(self.kernel_range)
if np.random.uniform() < self.opt['sinc_prob2']:
if kernel_size < 13:
omega_c = np.random.uniform(np.pi / 3, np.pi)
else:
omega_c = np.random.uniform(np.pi / 5, np.pi)
kernel2 = circular_lowpass_kernel(omega_c, kernel_size, pad_to=False)
else:
kernel2 = random_mixed_kernels(
self.kernel_list2,
self.kernel_prob2,
kernel_size,
self.blur_sigma2,
self.blur_sigma2, [-math.pi, math.pi],
self.betag_range2,
self.betap_range2,
noise_range=None)
# pad kernel
pad_size = (21 - kernel_size) // 2
kernel2 = np.pad(kernel2, ((pad_size, pad_size), (pad_size, pad_size)))
# ------------------------------------- the final sinc kernel ------------------------------------- #
if np.random.uniform() < self.opt['final_sinc_prob']:
kernel_size = random.choice(self.kernel_range)
omega_c = np.random.uniform(np.pi / 3, np.pi)
sinc_kernel = circular_lowpass_kernel(omega_c, kernel_size, pad_to=21)
sinc_kernel = torch.FloatTensor(sinc_kernel)
else:
sinc_kernel = self.pulse_tensor
# BGR to RGB, HWC to CHW, numpy to tensor
img_gt = img2tensor([img_gt], bgr2rgb=True, float32=True)[0]
kernel = torch.FloatTensor(kernel)
kernel2 = torch.FloatTensor(kernel2)
return_d = {'gt': img_gt, 'kernel1': kernel, 'kernel2': kernel2, 'sinc_kernel': sinc_kernel, 'gt_path': gt_path}
return return_d
def __len__(self):
return len(self.paths)
================================================
FILE: KDSR-GAN/kdsrgan/data/kdsrgan_paired_dataset.py
================================================
import os
from basicsr.data.data_util import paired_paths_from_folder, paired_paths_from_lmdb
from basicsr.data.transforms import augment, paired_random_crop
from basicsr.utils import FileClient, imfrombytes, img2tensor
from basicsr.utils.registry import DATASET_REGISTRY
from torch.utils import data as data
from torchvision.transforms.functional import normalize
@DATASET_REGISTRY.register()
class KDSRGANPairedDataset(data.Dataset):
"""Paired image dataset for image restoration.
Read LQ (Low Quality, e.g. LR (Low Resolution), blurry, noisy, etc) and GT image pairs.
There are three modes:
1. 'lmdb': Use lmdb files.
If opt['io_backend'] == lmdb.
2. 'meta_info': Use meta information file to generate paths.
If opt['io_backend'] != lmdb and opt['meta_info'] is not None.
3. 'folder': Scan folders to generate paths.
The rest.
Args:
opt (dict): Config for train datasets. It contains the following keys:
dataroot_gt (str): Data root path for gt.
dataroot_lq (str): Data root path for lq.
meta_info (str): Path for meta information file.
io_backend (dict): IO backend type and other kwarg.
filename_tmpl (str): Template for each filename. Note that the template excludes the file extension.
Default: '{}'.
gt_size (int): Cropped patched size for gt patches.
use_hflip (bool): Use horizontal flips.
use_rot (bool): Use rotation (use vertical flip and transposing h
and w for implementation).
scale (bool): Scale, which will be added automatically.
phase (str): 'train' or 'val'.
"""
def __init__(self, opt):
super(KDSRGANPairedDataset, self).__init__()
self.opt = opt
self.file_client = None
self.io_backend_opt = opt['io_backend']
# mean and std for normalizing the input images
self.mean = opt['mean'] if 'mean' in opt else None
self.std = opt['std'] if 'std' in opt else None
self.gt_folder, self.lq_folder = opt['dataroot_gt'], opt['dataroot_lq']
self.filename_tmpl = opt['filename_tmpl'] if 'filename_tmpl' in opt else '{}'
# file client (lmdb io backend)
if self.io_backend_opt['type'] == 'lmdb':
self.io_backend_opt['db_paths'] = [self.lq_folder, self.gt_folder]
self.io_backend_opt['client_keys'] = ['lq', 'gt']
self.paths = paired_paths_from_lmdb([self.lq_folder, self.gt_folder], ['lq', 'gt'])
elif 'meta_info' in self.opt and self.opt['meta_info'] is not None:
# disk backend with meta_info
# Each line in the meta_info describes the relative path to an image
with open(self.opt['meta_info']) as fin:
paths = [line.strip() for line in fin]
self.paths = []
for path in paths:
gt_path, lq_path = path.split(', ')
gt_path = os.path.join(self.gt_folder, gt_path)
lq_path = os.path.join(self.lq_folder, lq_path)
self.paths.append(dict([('gt_path', gt_path), ('lq_path', lq_path)]))
else:
# disk backend
# it will scan the whole folder to get meta info
# it will be time-consuming for folders with too many files. It is recommended using an extra meta txt file
self.paths = paired_paths_from_folder([self.lq_folder, self.gt_folder], ['lq', 'gt'], self.filename_tmpl)
def __getitem__(self, index):
if self.file_client is None:
self.file_client = FileClient(self.io_backend_opt.pop('type'), **self.io_backend_opt)
scale = self.opt['scale']
# Load gt and lq images. Dimension order: HWC; channel order: BGR;
# image range: [0, 1], float32.
gt_path = self.paths[index]['gt_path']
img_bytes = self.file_client.get(gt_path, 'gt')
img_gt = imfrombytes(img_bytes, float32=True)
lq_path = self.paths[index]['lq_path']
img_bytes = self.file_client.get(lq_path, 'lq')
img_lq = imfrombytes(img_bytes, float32=True)
# augmentation for training
if self.opt['phase'] == 'train':
gt_size = self.opt['gt_size']
# random crop
img_gt, img_lq = paired_random_crop(img_gt, img_lq, gt_size, scale, gt_path)
# flip, rotation
img_gt, img_lq = augment([img_gt, img_lq], self.opt['use_hflip'], self.opt['use_rot'])
# BGR to RGB, HWC to CHW, numpy to tensor
img_gt, img_lq = img2tensor([img_gt, img_lq], bgr2rgb=True, float32=True)
# normalize
if self.mean is not None or self.std is not None:
normalize(img_lq, self.mean, self.std, inplace=True)
normalize(img_gt, self.mean, self.std, inplace=True)
return {'lq': img_lq, 'gt': img_gt, 'lq_path': lq_path, 'gt_path': gt_path}
def __len__(self):
return len(self.paths)
================================================
FILE: KDSR-GAN/kdsrgan/losses/__init__.py
================================================
import importlib
from basicsr.utils import scandir
from os import path as osp
# automatically scan and import arch modules for registry
# scan all the files that end with '_arch.py' under the archs folder
arch_folder = osp.dirname(osp.abspath(__file__))
arch_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(arch_folder) if v.endswith('_loss.py')]
# import all the arch modules
_arch_modules = [importlib.import_module(f'kdsrgan.losses.{file_name}') for file_name in arch_filenames]
================================================
FILE: KDSR-GAN/kdsrgan/losses/my_loss.py
================================================
import torch
from torch import nn as nn
from torch.nn import functional as F
from basicsr.utils.registry import LOSS_REGISTRY
@LOSS_REGISTRY.register()
class KDLoss(nn.Module):
"""Knowledge distillation loss.
Args:
loss_weight (float): Loss weight for KD loss. Default: 1.0.
"""
def __init__(self, loss_weight=1.0, temperature = 0.15):
super(KDLoss, self).__init__()
self.loss_weight = loss_weight
self.temperature = temperature
def forward(self, T_fea, S_fea):
"""
Args:
T_fea (List): contain shape (N, L) vector of BlindSR_TA.
S_fea (List): contain shape (N, L) vector of BlindSR_ST.
weight (Tensor, optional): of shape (N, C, H, W). Element-wise weights. Default: None.
"""
loss_distill_dis = 0
for i in range(len(T_fea)):
student_distance = F.log_softmax(S_fea[i] / self.temperature, dim=1)
teacher_distance = F.softmax(T_fea[i].detach()/ self.temperature, dim=1)
loss_distill_dis += F.kl_div(
student_distance, teacher_distance, reduction='batchmean')
#loss_distill_abs += nn.L1Loss()(S_fea[i], T_fea[i].detach())
return self.loss_weight * loss_distill_dis
================================================
FILE: KDSR-GAN/kdsrgan/models/__init__.py
================================================
import importlib
from basicsr.utils import scandir
from os import path as osp
# automatically scan and import model modules for registry
# scan all the files that end with '_model.py' under the model folder
model_folder = osp.dirname(osp.abspath(__file__))
model_filenames = [osp.splitext(osp.basename(v))[0] for v in scandir(model_folder) if v.endswith('_model.py')]
# import all the model modules
_model_modules = [importlib.import_module(f'kdsrgan.models.{file_name}') for file_name in model_filenames]
================================================
FILE: KDSR-GAN/kdsrgan/models/kdsrgan_ST_model.py
================================================
import numpy as np
import random
import torch
from basicsr.data.degradations import random_add_gaussian_noise_pt, random_add_poisson_noise_pt
from basicsr.data.transforms import paired_random_crop
from basicsr.models.srgan_model import SRGANModel
from basicsr.utils import DiffJPEG, USMSharp
from basicsr.utils.img_process_util import filter2D
from basicsr.utils.registry import MODEL_REGISTRY
from collections import OrderedDict
from torch.nn import functional as F
from basicsr.archs import build_network
from basicsr.utils import get_root_logger
from basicsr.losses import build_loss
from torch import nn
@MODEL_REGISTRY.register()
class KDSRGANSTModel(SRGANModel):
"""
It mainly performs:
1. randomly synthesize LQ images in GPU tensors
2. optimize the networks with GAN training.
"""
def __init__(self, opt):
super(KDSRGANSTModel, self).__init__(opt)
self.jpeger = DiffJPEG(differentiable=False).cuda() # simulate JPEG compression artifacts
self.usm_sharpener = USMSharp().cuda() # do usm sharpening
self.queue_size = opt.get('queue_size', 180)
self.net_g_TA = build_network(opt['network_TA'])
self.net_g_TA = self.model_to_device(self.net_g_TA)
load_path = self.opt['path'].get('pretrain_network_TA', None)
if load_path is not None:
param_key = self.opt['path'].get('param_key_g', 'params')
self.load_network(self.net_g_TA, load_path, True, param_key)
self.net_g_TA.eval()
if self.opt['dist']:
self.model_Est = self.net_g.module.E_st
self.model_Eta = self.net_g_TA.module.E
else:
self.model_Est = self.net_g.E_st
self.model_Eta = self.net_g_TA.E
self.pixel_unshuffle = nn.PixelUnshuffle(opt["scale"])
if self.is_train:
self.encoder_iter = opt["train"]["encoder_iter"]
self.lr_encoder = opt["train"]["lr_encoder"]
self.lr_sr = opt["train"]["lr_sr"]
self.gamma_encoder = opt["train"]["gamma_encoder"]
self.gamma_sr = opt["train"]["gamma_sr"]
self.lr_decay_encoder = opt["train"]["lr_decay_encoder"]
self.lr_decay_sr = opt["train"]["lr_decay_sr"]
def init_training_settings(self):
train_opt = self.opt['train']
if train_opt.get('kd_opt'):
self.cri_kd = build_loss(train_opt['kd_opt']).to(self.device)
else:
self.cri_kd = None
super(KDSRGANSTModel, self).init_training_settings()
@torch.no_grad()
def _dequeue_and_enqueue(self):
"""It is the training pair pool for increasing the diversity in a batch.
Batch processing limits the diversity of synthetic degradations in a batch. For example, samples in a
batch could not have different resize scaling factors. Therefore, we employ this training pair pool
to increase the degradation diversity in a batch.
"""
# initialize
b, c, h, w = self.lq.size()
if not hasattr(self, 'queue_lr'):
assert self.queue_size % b == 0, f'queue size {self.queue_size} should be divisible by batch size {b}'
self.queue_lr = torch.zeros(self.queue_size, c, h, w).cuda()
_, c, h, w = self.gt.size()
self.queue_gt = torch.zeros(self.queue_size, c, h, w).cuda()
self.queue_ptr = 0
if self.queue_ptr == self.queue_size: # the pool is full
# do dequeue and enqueue
# shuffle
idx = torch.randperm(self.queue_size)
self.queue_lr = self.queue_lr[idx]
self.queue_gt = self.queue_gt[idx]
# get first b samples
lq_dequeue = self.queue_lr[0:b, :, :, :].clone()
gt_dequeue = self.queue_gt[0:b, :, :, :].clone()
# update the queue
self.queue_lr[0:b, :, :, :] = self.lq.clone()
self.queue_gt[0:b, :, :, :] = self.gt.clone()
self.lq = lq_dequeue
self.gt = gt_dequeue
else:
# only do enqueue
self.queue_lr[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.lq.clone()
self.queue_gt[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.gt.clone()
self.queue_ptr = self.queue_ptr + b
@torch.no_grad()
def feed_data(self, data):
"""Accept data from dataloader, and then add two-order degradations to obtain LQ images.
"""
if self.is_train and self.opt.get('high_order_degradation', True):
# training data synthesis
self.gt = data['gt'].to(self.device)
self.gt_usm = self.usm_sharpener(self.gt)
self.kernel1 = data['kernel1'].to(self.device)
self.kernel2 = data['kernel2'].to(self.device)
self.sinc_kernel = data['sinc_kernel'].to(self.device)
ori_h, ori_w = self.gt.size()[2:4]
# ----------------------- The first degradation process ----------------------- #
# blur
out = filter2D(self.gt, self.kernel1)
# random resize
updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob'])[0]
if updown_type == 'up':
scale = np.random.uniform(1, self.opt['resize_range'][1])
elif updown_type == 'down':
scale = np.random.uniform(self.opt['resize_range'][0], 1)
else:
scale = 1
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(out, scale_factor=scale, mode=mode)
# add noise
gray_noise_prob = self.opt['gray_noise_prob']
if np.random.uniform() < self.opt['gaussian_noise_prob']:
out = random_add_gaussian_noise_pt(
out, sigma_range=self.opt['noise_range'], clip=True, rounds=False, gray_prob=gray_noise_prob)
else:
out = random_add_poisson_noise_pt(
out,
scale_range=self.opt['poisson_scale_range'],
gray_prob=gray_noise_prob,
clip=True,
rounds=False)
# JPEG compression
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range'])
out = torch.clamp(out, 0, 1) # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts
out = self.jpeger(out, quality=jpeg_p)
# ----------------------- The second degradation process ----------------------- #
# blur
if np.random.uniform() < self.opt['second_blur_prob']:
out = filter2D(out, self.kernel2)
# random resize
updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob2'])[0]
if updown_type == 'up':
scale = np.random.uniform(1, self.opt['resize_range2'][1])
elif updown_type == 'down':
scale = np.random.uniform(self.opt['resize_range2'][0], 1)
else:
scale = 1
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(
out, size=(int(ori_h / self.opt['scale'] * scale), int(ori_w / self.opt['scale'] * scale)), mode=mode)
# add noise
gray_noise_prob = self.opt['gray_noise_prob2']
if np.random.uniform() < self.opt['gaussian_noise_prob2']:
out = random_add_gaussian_noise_pt(
out, sigma_range=self.opt['noise_range2'], clip=True, rounds=False, gray_prob=gray_noise_prob)
else:
out = random_add_poisson_noise_pt(
out,
scale_range=self.opt['poisson_scale_range2'],
gray_prob=gray_noise_prob,
clip=True,
rounds=False)
# JPEG compression + the final sinc filter
# We also need to resize images to desired sizes. We group [resize back + sinc filter] together
# as one operation.
# We consider two orders:
# 1. [resize back + sinc filter] + JPEG compression
# 2. JPEG compression + [resize back + sinc filter]
# Empirically, we find other combinations (sinc + JPEG + Resize) will introduce twisted lines.
if np.random.uniform() < 0.5:
# resize back + the final sinc filter
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
out = filter2D(out, self.sinc_kernel)
# JPEG compression
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
out = torch.clamp(out, 0, 1)
out = self.jpeger(out, quality=jpeg_p)
else:
# JPEG compression
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
out = torch.clamp(out, 0, 1)
out = self.jpeger(out, quality=jpeg_p)
# resize back + the final sinc filter
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
out = filter2D(out, self.sinc_kernel)
# clamp and round
self.lq = torch.clamp((out * 255.0).round(), 0, 255) / 255.
# random crop
gt_size = self.opt['gt_size']
(self.gt, self.gt_usm), self.lq = paired_random_crop([self.gt, self.gt_usm], self.lq, gt_size,
self.opt['scale'])
# training pair pool
self._dequeue_and_enqueue()
# sharpen self.gt again, as we have changed the self.gt with self._dequeue_and_enqueue
self.gt_usm = self.usm_sharpener(self.gt)
self.lq = self.lq.contiguous() # for the warning: grad and param do not obey the gradient layout contract
else:
# for paired training or validation
self.lq = data['lq'].to(self.device)
if 'gt' in data:
self.gt = data['gt'].to(self.device)
self.gt_usm = self.usm_sharpener(self.gt)
def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
# do not use the synthetic process during validation
self.is_train = False
super(KDSRGANSTModel, self).nondist_validation(dataloader, current_iter, tb_logger, save_img)
self.is_train = True
def test(self):
if hasattr(self, 'net_g_ema'):
self.net_g_ema.eval()
with torch.no_grad():
self.output = self.net_g_ema(self.lq)
else:
self.net_g.eval()
with torch.no_grad():
self.output = self.net_g(self.lq)
self.net_g.train()
def optimize_parameters(self, current_iter):
lr = self.lr_sr * (self.gamma_sr ** ((current_iter ) // self.lr_decay_sr))
for param_group in self.optimizer_g.param_groups:
param_group['lr'] = lr
l1_gt = self.gt_usm
percep_gt = self.gt_usm
gan_gt = self.gt_usm
if self.opt['l1_gt_usm'] is False:
l1_gt = self.gt
if self.opt['percep_gt_usm'] is False:
percep_gt = self.gt
if self.opt['gan_gt_usm'] is False:
gan_gt = self.gt
hr2 = self.pixel_unshuffle(l1_gt)
_, T_fea = self.model_Eta(torch.cat([self.lq,hr2],dim=1))
# optimize net_g
for p in self.net_d.parameters():
p.requires_grad = False
self.optimizer_g.zero_grad()
self.output,S_fea = self.net_g(self.lq)
l_g_total = 0
loss_dict = OrderedDict()
if (current_iter % self.net_d_iters == 0 and current_iter > self.net_d_init_iters):
# pixel loss
if self.cri_pix:
l_g_pix = self.cri_pix(self.output, self.gt)
l_g_total += l_g_pix
loss_dict['l_g_pix'] = l_g_pix
if self.cri_kd:
l_kd = self.cri_kd(T_fea, S_fea)
l_g_total += l_kd
loss_dict['l_kd'] = l_kd
# perceptual loss
if self.cri_perceptual:
l_g_percep, l_g_style = self.cri_perceptual(self.output, self.gt)
if l_g_percep is not None:
l_g_total += l_g_percep
loss_dict['l_g_percep'] = l_g_percep
if l_g_style is not None:
l_g_total += l_g_style
loss_dict['l_g_style'] = l_g_style
# gan loss
fake_g_pred = self.net_d(self.output)
l_g_gan = self.cri_gan(fake_g_pred, True, is_disc=False)
l_g_total += l_g_gan
loss_dict['l_g_gan'] = l_g_gan
l_g_total.backward()
self.optimizer_g.step()
# optimize net_d
for p in self.net_d.parameters():
p.requires_grad = True
self.optimizer_d.zero_grad()
# real
real_d_pred = self.net_d(self.gt)
l_d_real = self.cri_gan(real_d_pred, True, is_disc=True)
loss_dict['l_d_real'] = l_d_real
loss_dict['out_d_real'] = torch.mean(real_d_pred.detach())
l_d_real.backward()
# fake
fake_d_pred = self.net_d(self.output.detach().clone()) # clone for pt1.9
l_d_fake = self.cri_gan(fake_d_pred, False, is_disc=True)
loss_dict['l_d_fake'] = l_d_fake
loss_dict['out_d_fake'] = torch.mean(fake_d_pred.detach())
l_d_fake.backward()
self.optimizer_d.step()
if self.ema_decay > 0:
self.model_ema(decay=self.ema_decay)
self.log_dict = self.reduce_loss_dict(loss_dict)
================================================
FILE: KDSR-GAN/kdsrgan/models/kdsrgan_TA_model.py
================================================
import numpy as np
import random
import torch
from basicsr.data.degradations import random_add_gaussian_noise_pt, random_add_poisson_noise_pt
from basicsr.data.transforms import paired_random_crop
from basicsr.models.srgan_model import SRGANModel
from basicsr.utils import DiffJPEG, USMSharp
from basicsr.utils.img_process_util import filter2D
from basicsr.utils.registry import MODEL_REGISTRY
from collections import OrderedDict
from torch.nn import functional as F
@MODEL_REGISTRY.register()
class KDSRGANTAModel(SRGANModel):
"""
It mainly performs:
1. randomly synthesize LQ images in GPU tensors
2. optimize the networks with GAN training.
"""
def __init__(self, opt):
super(KDSRGANTAModel, self).__init__(opt)
self.jpeger = DiffJPEG(differentiable=False).cuda() # simulate JPEG compression artifacts
self.usm_sharpener = USMSharp().cuda() # do usm sharpening
self.queue_size = opt.get('queue_size', 180)
@torch.no_grad()
def _dequeue_and_enqueue(self):
"""It is the training pair pool for increasing the diversity in a batch.
Batch processing limits the diversity of synthetic degradations in a batch. For example, samples in a
batch could not have different resize scaling factors. Therefore, we employ this training pair pool
to increase the degradation diversity in a batch.
"""
# initialize
b, c, h, w = self.lq.size()
if not hasattr(self, 'queue_lr'):
assert self.queue_size % b == 0, f'queue size {self.queue_size} should be divisible by batch size {b}'
self.queue_lr = torch.zeros(self.queue_size, c, h, w).cuda()
_, c, h, w = self.gt.size()
self.queue_gt = torch.zeros(self.queue_size, c, h, w).cuda()
self.queue_ptr = 0
if self.queue_ptr == self.queue_size: # the pool is full
# do dequeue and enqueue
# shuffle
idx = torch.randperm(self.queue_size)
self.queue_lr = self.queue_lr[idx]
self.queue_gt = self.queue_gt[idx]
# get first b samples
lq_dequeue = self.queue_lr[0:b, :, :, :].clone()
gt_dequeue = self.queue_gt[0:b, :, :, :].clone()
# update the queue
self.queue_lr[0:b, :, :, :] = self.lq.clone()
self.queue_gt[0:b, :, :, :] = self.gt.clone()
self.lq = lq_dequeue
self.gt = gt_dequeue
else:
# only do enqueue
self.queue_lr[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.lq.clone()
self.queue_gt[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.gt.clone()
self.queue_ptr = self.queue_ptr + b
@torch.no_grad()
def feed_data(self, data):
"""Accept data from dataloader, and then add two-order degradations to obtain LQ images.
"""
if self.is_train and self.opt.get('high_order_degradation', True):
# training data synthesis
self.gt = data['gt'].to(self.device)
self.gt_usm = self.usm_sharpener(self.gt)
self.kernel1 = data['kernel1'].to(self.device)
self.kernel2 = data['kernel2'].to(self.device)
self.sinc_kernel = data['sinc_kernel'].to(self.device)
ori_h, ori_w = self.gt.size()[2:4]
# ----------------------- The first degradation process ----------------------- #
# blur
out = filter2D(self.gt_usm, self.kernel1)
# random resize
updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob'])[0]
if updown_type == 'up':
scale = np.random.uniform(1, self.opt['resize_range'][1])
elif updown_type == 'down':
scale = np.random.uniform(self.opt['resize_range'][0], 1)
else:
scale = 1
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(out, scale_factor=scale, mode=mode)
# add noise
gray_noise_prob = self.opt['gray_noise_prob']
if np.random.uniform() < self.opt['gaussian_noise_prob']:
out = random_add_gaussian_noise_pt(
out, sigma_range=self.opt['noise_range'], clip=True, rounds=False, gray_prob=gray_noise_prob)
else:
out = random_add_poisson_noise_pt(
out,
scale_range=self.opt['poisson_scale_range'],
gray_prob=gray_noise_prob,
clip=True,
rounds=False)
# JPEG compression
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range'])
out = torch.clamp(out, 0, 1) # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts
out = self.jpeger(out, quality=jpeg_p)
# ----------------------- The second degradation process ----------------------- #
# blur
if np.random.uniform() < self.opt['second_blur_prob']:
out = filter2D(out, self.kernel2)
# random resize
updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob2'])[0]
if updown_type == 'up':
scale = np.random.uniform(1, self.opt['resize_range2'][1])
elif updown_type == 'down':
scale = np.random.uniform(self.opt['resize_range2'][0], 1)
else:
scale = 1
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(
out, size=(int(ori_h / self.opt['scale'] * scale), int(ori_w / self.opt['scale'] * scale)), mode=mode)
# add noise
gray_noise_prob = self.opt['gray_noise_prob2']
if np.random.uniform() < self.opt['gaussian_noise_prob2']:
out = random_add_gaussian_noise_pt(
out, sigma_range=self.opt['noise_range2'], clip=True, rounds=False, gray_prob=gray_noise_prob)
else:
out = random_add_poisson_noise_pt(
out,
scale_range=self.opt['poisson_scale_range2'],
gray_prob=gray_noise_prob,
clip=True,
rounds=False)
# JPEG compression + the final sinc filter
# We also need to resize images to desired sizes. We group [resize back + sinc filter] together
# as one operation.
# We consider two orders:
# 1. [resize back + sinc filter] + JPEG compression
# 2. JPEG compression + [resize back + sinc filter]
# Empirically, we find other combinations (sinc + JPEG + Resize) will introduce twisted lines.
if np.random.uniform() < 0.5:
# resize back + the final sinc filter
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
out = filter2D(out, self.sinc_kernel)
# JPEG compression
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
out = torch.clamp(out, 0, 1)
out = self.jpeger(out, quality=jpeg_p)
else:
# JPEG compression
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
out = torch.clamp(out, 0, 1)
out = self.jpeger(out, quality=jpeg_p)
# resize back + the final sinc filter
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
out = filter2D(out, self.sinc_kernel)
# clamp and round
self.lq = torch.clamp((out * 255.0).round(), 0, 255) / 255.
# random crop
gt_size = self.opt['gt_size']
(self.gt, self.gt_usm), self.lq = paired_random_crop([self.gt, self.gt_usm], self.lq, gt_size,
self.opt['scale'])
# training pair pool
self._dequeue_and_enqueue()
# sharpen self.gt again, as we have changed the self.gt with self._dequeue_and_enqueue
self.gt_usm = self.usm_sharpener(self.gt)
self.lq = self.lq.contiguous() # for the warning: grad and param do not obey the gradient layout contract
else:
# for paired training or validation
self.lq = data['lq'].to(self.device)
if 'gt' in data:
self.gt = data['gt'].to(self.device)
self.gt_usm = self.usm_sharpener(self.gt)
def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
# do not use the synthetic process during validation
self.is_train = False
super(KDSRGANTAModel, self).nondist_validation(dataloader, current_iter, tb_logger, save_img)
self.is_train = True
def optimize_parameters(self, current_iter):
# optimize net_g
for p in self.net_d.parameters():
p.requires_grad = False
self.optimizer_g.zero_grad()
self.output = self.net_g(self.lq)
l_g_total = 0
loss_dict = OrderedDict()
if (current_iter % self.net_d_iters == 0 and current_iter > self.net_d_init_iters):
# pixel loss
if self.cri_pix:
l_g_pix = self.cri_pix(self.output, self.gt)
l_g_total += l_g_pix
loss_dict['l_g_pix'] = l_g_pix
# perceptual loss
if self.cri_perceptual:
l_g_percep, l_g_style = self.cri_perceptual(self.output, self.gt)
if l_g_percep is not None:
l_g_total += l_g_percep
loss_dict['l_g_percep'] = l_g_percep
if l_g_style is not None:
l_g_total += l_g_style
loss_dict['l_g_style'] = l_g_style
# gan loss
fake_g_pred = self.net_d(self.output)
l_g_gan = self.cri_gan(fake_g_pred, True, is_disc=False)
l_g_total += l_g_gan
loss_dict['l_g_gan'] = l_g_gan
l_g_total.backward()
self.optimizer_g.step()
# optimize net_d
for p in self.net_d.parameters():
p.requires_grad = True
self.optimizer_d.zero_grad()
# real
real_d_pred = self.net_d(self.gt)
l_d_real = self.cri_gan(real_d_pred, True, is_disc=True)
loss_dict['l_d_real'] = l_d_real
loss_dict['out_d_real'] = torch.mean(real_d_pred.detach())
l_d_real.backward()
# fake
fake_d_pred = self.net_d(self.output.detach().clone()) # clone for pt1.9
l_d_fake = self.cri_gan(fake_d_pred, False, is_disc=True)
loss_dict['l_d_fake'] = l_d_fake
loss_dict['out_d_fake'] = torch.mean(fake_d_pred.detach())
l_d_fake.backward()
self.optimizer_d.step()
if self.ema_decay > 0:
self.model_ema(decay=self.ema_decay)
self.log_dict = self.reduce_loss_dict(loss_dict)
================================================
FILE: KDSR-GAN/kdsrgan/models/kdsrnet_ST_model.py
================================================
import numpy as np
import random
import torch
from basicsr.data.degradations import random_add_gaussian_noise_pt, random_add_poisson_noise_pt
from basicsr.data.transforms import paired_random_crop
from basicsr.models.sr_model import SRModel
from basicsr.utils import DiffJPEG, USMSharp
from basicsr.utils.img_process_util import filter2D
from basicsr.utils.registry import MODEL_REGISTRY
from torch.nn import functional as F
from collections import OrderedDict
from basicsr.archs import build_network
from basicsr.utils import get_root_logger
from basicsr.losses import build_loss
from torch import nn
from basicsr.models import lr_scheduler as lr_scheduler
@MODEL_REGISTRY.register()
class KDSRNetSTModel(SRModel):
"""
It is trained without GAN losses.
It mainly performs:
1. randomly synthesize LQ images in GPU tensors
2. optimize the networks with GAN training.
"""
def __init__(self, opt):
super(KDSRNetSTModel, self).__init__(opt)
self.jpeger = DiffJPEG(differentiable=False).cuda() # simulate JPEG compression artifacts
self.usm_sharpener = USMSharp().cuda() # do usm sharpening
self.queue_size = opt.get('queue_size', 180)
self.net_g_TA = build_network(opt['network_TA'])
self.net_g_TA = self.model_to_device(self.net_g_TA)
# load pretrained models
load_path = self.opt['path'].get('pretrain_network_TA', None)
if load_path is not None:
param_key = self.opt['path'].get('param_key_g', 'params')
self.load_network(self.net_g_TA, load_path, True, param_key)
self.net_g_TA.eval()
if self.opt['dist']:
self.model_Est = self.net_g.module.E_st
self.model_Eta = self.net_g_TA.module.E
else:
self.model_Est = self.net_g.E_st
self.model_Eta = self.net_g_TA.E
self.pixel_unshuffle = nn.PixelUnshuffle(opt["scale"])
self.encoder_iter = opt["train"]["encoder_iter"]
self.lr_encoder = opt["train"]["lr_encoder"]
self.lr_sr = opt["train"]["lr_sr"]
self.gamma_encoder = opt["train"]["gamma_encoder"]
self.gamma_sr = opt["train"]["gamma_sr"]
self.lr_decay_encoder = opt["train"]["lr_decay_encoder"]
self.lr_decay_sr = opt["train"]["lr_decay_sr"]
def setup_optimizers(self):
train_opt = self.opt['train']
optim_params = []
for k, v in self.net_g.named_parameters():
if v.requires_grad:
optim_params.append(v)
else:
logger = get_root_logger()
logger.warning(f'Params {k} will not be optimized in the second stage.')
optim_type = train_opt['optim_g'].pop('type')
self.optimizer_g = self.get_optimizer(optim_type, optim_params, **train_opt['optim_g'])
self.optimizers.append(self.optimizer_g)
def setup_schedulers(self):
"""Set up schedulers."""
train_opt = self.opt['train']
scheduler_type = train_opt['scheduler'].pop('type')
if scheduler_type in ['MultiStepLR', 'MultiStepRestartLR']:
self.schedulers.append(lr_scheduler.MultiStepRestartLR(self.optimizer_g, **train_opt['scheduler']))
elif scheduler_type == 'CosineAnnealingRestartLR':
self.schedulers.append(lr_scheduler.CosineAnnealingRestartLR(self.optimizer_g, **train_opt['scheduler']))
else:
raise NotImplementedError(f'Scheduler {scheduler_type} is not implemented yet.')
def init_training_settings(self):
self.net_g.train()
train_opt = self.opt['train']
self.ema_decay = train_opt.get('ema_decay', 0)
if self.ema_decay > 0:
logger = get_root_logger()
logger.info(f'Use Exponential Moving Average with decay: {self.ema_decay}')
# define network net_g with Exponential Moving Average (EMA)
# net_g_ema is used only for testing on one GPU and saving
# There is no need to wrap with DistributedDataParallel
self.net_g_ema = build_network(self.opt['network_g']).to(self.device)
# load pretrained model
load_path = self.opt['path'].get('pretrain_network_g', None)
if load_path is not None:
self.load_network(self.net_g_ema, load_path, self.opt['path'].get('strict_load_g', True), 'params_ema')
else:
self.model_ema(0) # copy net_g weight
self.net_g_ema.eval()
# define losses
if train_opt.get('pixel_opt'):
self.cri_pix = build_loss(train_opt['pixel_opt']).to(self.device)
else:
self.cri_pix = None
if train_opt.get('perceptual_opt'):
self.cri_perceptual = build_loss(train_opt['perceptual_opt']).to(self.device)
else:
self.cri_perceptual = None
if train_opt.get('kd_opt'):
self.cri_kd = build_loss(train_opt['kd_opt']).to(self.device)
else:
self.cri_kd = None
if self.cri_pix is None and self.cri_perceptual is None:
raise ValueError('Both pixel and perceptual losses are None.')
# set up optimizers and schedulers
self.setup_optimizers()
self.setup_schedulers()
@torch.no_grad()
def _dequeue_and_enqueue(self):
"""It is the training pair pool for increasing the diversity in a batch.
Batch processing limits the diversity of synthetic degradations in a batch. For example, samples in a
batch could not have different resize scaling factors. Therefore, we employ this training pair pool
to increase the degradation diversity in a batch.
"""
# initialize
b, c, h, w = self.lq.size()
if not hasattr(self, 'queue_lr'):
assert self.queue_size % b == 0, f'queue size {self.queue_size} should be divisible by batch size {b}'
self.queue_lr = torch.zeros(self.queue_size, c, h, w).cuda()
_, c, h, w = self.gt.size()
self.queue_gt = torch.zeros(self.queue_size, c, h, w).cuda()
self.queue_ptr = 0
if self.queue_ptr == self.queue_size: # the pool is full
# do dequeue and enqueue
# shuffle
idx = torch.randperm(self.queue_size)
self.queue_lr = self.queue_lr[idx]
self.queue_gt = self.queue_gt[idx]
# get first b samples
lq_dequeue = self.queue_lr[0:b, :, :, :].clone()
gt_dequeue = self.queue_gt[0:b, :, :, :].clone()
# update the queue
self.queue_lr[0:b, :, :, :] = self.lq.clone()
self.queue_gt[0:b, :, :, :] = self.gt.clone()
self.lq = lq_dequeue
self.gt = gt_dequeue
else:
# only do enqueue
self.queue_lr[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.lq.clone()
self.queue_gt[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.gt.clone()
self.queue_ptr = self.queue_ptr + b
@torch.no_grad()
def feed_data(self, data):
"""Accept data from dataloader, and then add two-order degradations to obtain LQ images.
"""
if self.is_train and self.opt.get('high_order_degradation', True):
# training data synthesis
self.gt = data['gt'].to(self.device)
# USM sharpen the GT images
if self.opt['gt_usm'] is True:
self.gt = self.usm_sharpener(self.gt)
self.kernel1 = data['kernel1'].to(self.device)
self.kernel2 = data['kernel2'].to(self.device)
self.sinc_kernel = data['sinc_kernel'].to(self.device)
ori_h, ori_w = self.gt.size()[2:4]
# ----------------------- The first degradation process ----------------------- #
# blur
out = filter2D(self.gt, self.kernel1)
# random resize
updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob'])[0]
if updown_type == 'up':
scale = np.random.uniform(1, self.opt['resize_range'][1])
elif updown_type == 'down':
scale = np.random.uniform(self.opt['resize_range'][0], 1)
else:
scale = 1
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(out, scale_factor=scale, mode=mode)
# add noise
gray_noise_prob = self.opt['gray_noise_prob']
if np.random.uniform() < self.opt['gaussian_noise_prob']:
out = random_add_gaussian_noise_pt(
out, sigma_range=self.opt['noise_range'], clip=True, rounds=False, gray_prob=gray_noise_prob)
else:
out = random_add_poisson_noise_pt(
out,
scale_range=self.opt['poisson_scale_range'],
gray_prob=gray_noise_prob,
clip=True,
rounds=False)
# JPEG compression
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range'])
out = torch.clamp(out, 0, 1) # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts
out = self.jpeger(out, quality=jpeg_p)
# ----------------------- The second degradation process ----------------------- #
# blur
if np.random.uniform() < self.opt['second_blur_prob']:
out = filter2D(out, self.kernel2)
# random resize
updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob2'])[0]
if updown_type == 'up':
scale = np.random.uniform(1, self.opt['resize_range2'][1])
elif updown_type == 'down':
scale = np.random.uniform(self.opt['resize_range2'][0], 1)
else:
scale = 1
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(
out, size=(int(ori_h / self.opt['scale'] * scale), int(ori_w / self.opt['scale'] * scale)), mode=mode)
# add noise
gray_noise_prob = self.opt['gray_noise_prob2']
if np.random.uniform() < self.opt['gaussian_noise_prob2']:
out = random_add_gaussian_noise_pt(
out, sigma_range=self.opt['noise_range2'], clip=True, rounds=False, gray_prob=gray_noise_prob)
else:
out = random_add_poisson_noise_pt(
out,
scale_range=self.opt['poisson_scale_range2'],
gray_prob=gray_noise_prob,
clip=True,
rounds=False)
# JPEG compression + the final sinc filter
# We also need to resize images to desired sizes. We group [resize back + sinc filter] together
# as one operation.
# We consider two orders:
# 1. [resize back + sinc filter] + JPEG compression
# 2. JPEG compression + [resize back + sinc filter]
# Empirically, we find other combinations (sinc + JPEG + Resize) will introduce twisted lines.
if np.random.uniform() < 0.5:
# resize back + the final sinc filter
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
out = filter2D(out, self.sinc_kernel)
# JPEG compression
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
out = torch.clamp(out, 0, 1)
out = self.jpeger(out, quality=jpeg_p)
else:
# JPEG compression
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
out = torch.clamp(out, 0, 1)
out = self.jpeger(out, quality=jpeg_p)
# resize back + the final sinc filter
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
out = filter2D(out, self.sinc_kernel)
# clamp and round
self.lq = torch.clamp((out * 255.0).round(), 0, 255) / 255.
# random crop
gt_size = self.opt['gt_size']
self.gt, self.lq = paired_random_crop(self.gt, self.lq, gt_size, self.opt['scale'])
# training pair pool
self._dequeue_and_enqueue()
self.lq = self.lq.contiguous() # for the warning: grad and param do not obey the gradient layout contract
else:
# for paired training or validation
self.lq = data['lq'].to(self.device)
if 'gt' in data:
self.gt = data['gt'].to(self.device)
self.gt_usm = self.usm_sharpener(self.gt)
def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
# do not use the synthetic process during validation
self.is_train = False
super(KDSRNetSTModel, self).nondist_validation(dataloader, current_iter, tb_logger, save_img)
self.is_train = True
def test(self):
if hasattr(self, 'net_g_ema'):
self.net_g_ema.eval()
with torch.no_grad():
self.output = self.net_g_ema(self.lq)
else:
self.net_g.eval()
with torch.no_grad():
self.output = self.net_g(self.lq)
self.net_g.train()
def optimize_parameters(self, current_iter):
lr = self.lr_sr * (self.gamma_sr ** ((current_iter ) // self.lr_decay_sr))
for param_group in self.optimizer_g.param_groups:
param_group['lr'] = lr
l_total = 0
loss_dict = OrderedDict()
hr2 = self.pixel_unshuffle(self.gt)
_, T_fea = self.model_Eta(torch.cat([self.lq,hr2],dim=1))
self.optimizer_g.zero_grad()
self.output, S_fea = self.net_g(self.lq)
l_pix = self.cri_pix(self.output, self.gt)
l_total += l_pix
loss_dict['l_pix'] = l_pix
l_kd = self.cri_kd(T_fea, S_fea)
l_total += l_kd
loss_dict['l_kd'] = l_kd
l_total.backward()
self.optimizer_g.step()
self.log_dict = self.reduce_loss_dict(loss_dict)
if self.ema_decay > 0:
self.model_ema(decay=self.ema_decay)
================================================
FILE: KDSR-GAN/kdsrgan/models/kdsrnet_TA_model.py
================================================
import numpy as np
import random
import torch
from basicsr.data.degradations import random_add_gaussian_noise_pt, random_add_poisson_noise_pt
from basicsr.data.transforms import paired_random_crop
from basicsr.models.sr_model import SRModel
from basicsr.utils import DiffJPEG, USMSharp
from basicsr.utils.img_process_util import filter2D
from basicsr.utils.registry import MODEL_REGISTRY
from torch.nn import functional as F
from collections import OrderedDict
@MODEL_REGISTRY.register()
class KDSRNetTAModel(SRModel):
"""
It is trained without GAN losses.
It mainly performs:
1. randomly synthesize LQ images in GPU tensors
2. optimize the networks with GAN training.
"""
def __init__(self, opt):
super(KDSRNetTAModel, self).__init__(opt)
self.jpeger = DiffJPEG(differentiable=False).cuda() # simulate JPEG compression artifacts
self.usm_sharpener = USMSharp().cuda() # do usm sharpening
self.queue_size = opt.get('queue_size', 180)
@torch.no_grad()
def _dequeue_and_enqueue(self):
"""It is the training pair pool for increasing the diversity in a batch.
Batch processing limits the diversity of synthetic degradations in a batch. For example, samples in a
batch could not have different resize scaling factors. Therefore, we employ this training pair pool
to increase the degradation diversity in a batch.
"""
# initialize
b, c, h, w = self.lq.size()
if not hasattr(self, 'queue_lr'):
assert self.queue_size % b == 0, f'queue size {self.queue_size} should be divisible by batch size {b}'
self.queue_lr = torch.zeros(self.queue_size, c, h, w).cuda()
_, c, h, w = self.gt.size()
self.queue_gt = torch.zeros(self.queue_size, c, h, w).cuda()
self.queue_ptr = 0
if self.queue_ptr == self.queue_size: # the pool is full
# do dequeue and enqueue
# shuffle
idx = torch.randperm(self.queue_size)
self.queue_lr = self.queue_lr[idx]
self.queue_gt = self.queue_gt[idx]
# get first b samples
lq_dequeue = self.queue_lr[0:b, :, :, :].clone()
gt_dequeue = self.queue_gt[0:b, :, :, :].clone()
# update the queue
self.queue_lr[0:b, :, :, :] = self.lq.clone()
self.queue_gt[0:b, :, :, :] = self.gt.clone()
self.lq = lq_dequeue
self.gt = gt_dequeue
else:
# only do enqueue
self.queue_lr[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.lq.clone()
self.queue_gt[self.queue_ptr:self.queue_ptr + b, :, :, :] = self.gt.clone()
self.queue_ptr = self.queue_ptr + b
@torch.no_grad()
def feed_data(self, data):
"""Accept data from dataloader, and then add two-order degradations to obtain LQ images.
"""
if self.is_train and self.opt.get('high_order_degradation', True):
# training data synthesis
self.gt = data['gt'].to(self.device)
# USM sharpen the GT images
if self.opt['gt_usm'] is True:
self.gt = self.usm_sharpener(self.gt)
self.kernel1 = data['kernel1'].to(self.device)
self.kernel2 = data['kernel2'].to(self.device)
self.sinc_kernel = data['sinc_kernel'].to(self.device)
ori_h, ori_w = self.gt.size()[2:4]
# ----------------------- The first degradation process ----------------------- #
# blur
out = filter2D(self.gt, self.kernel1)
# random resize
updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob'])[0]
if updown_type == 'up':
scale = np.random.uniform(1, self.opt['resize_range'][1])
elif updown_type == 'down':
scale = np.random.uniform(self.opt['resize_range'][0], 1)
else:
scale = 1
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(out, scale_factor=scale, mode=mode)
# add noise
gray_noise_prob = self.opt['gray_noise_prob']
if np.random.uniform() < self.opt['gaussian_noise_prob']:
out = random_add_gaussian_noise_pt(
out, sigma_range=self.opt['noise_range'], clip=True, rounds=False, gray_prob=gray_noise_prob)
else:
out = random_add_poisson_noise_pt(
out,
scale_range=self.opt['poisson_scale_range'],
gray_prob=gray_noise_prob,
clip=True,
rounds=False)
# JPEG compression
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range'])
out = torch.clamp(out, 0, 1) # clamp to [0, 1], otherwise JPEGer will result in unpleasant artifacts
out = self.jpeger(out, quality=jpeg_p)
# ----------------------- The second degradation process ----------------------- #
# blur
if np.random.uniform() < self.opt['second_blur_prob']:
out = filter2D(out, self.kernel2)
# random resize
updown_type = random.choices(['up', 'down', 'keep'], self.opt['resize_prob2'])[0]
if updown_type == 'up':
scale = np.random.uniform(1, self.opt['resize_range2'][1])
elif updown_type == 'down':
scale = np.random.uniform(self.opt['resize_range2'][0], 1)
else:
scale = 1
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(
out, size=(int(ori_h / self.opt['scale'] * scale), int(ori_w / self.opt['scale'] * scale)), mode=mode)
# add noise
gray_noise_prob = self.opt['gray_noise_prob2']
if np.random.uniform() < self.opt['gaussian_noise_prob2']:
out = random_add_gaussian_noise_pt(
out, sigma_range=self.opt['noise_range2'], clip=True, rounds=False, gray_prob=gray_noise_prob)
else:
out = random_add_poisson_noise_pt(
out,
scale_range=self.opt['poisson_scale_range2'],
gray_prob=gray_noise_prob,
clip=True,
rounds=False)
# JPEG compression + the final sinc filter
# We also need to resize images to desired sizes. We group [resize back + sinc filter] together
# as one operation.
# We consider two orders:
# 1. [resize back + sinc filter] + JPEG compression
# 2. JPEG compression + [resize back + sinc filter]
# Empirically, we find other combinations (sinc + JPEG + Resize) will introduce twisted lines.
if np.random.uniform() < 0.5:
# resize back + the final sinc filter
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
out = filter2D(out, self.sinc_kernel)
# JPEG compression
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
out = torch.clamp(out, 0, 1)
out = self.jpeger(out, quality=jpeg_p)
else:
# JPEG compression
jpeg_p = out.new_zeros(out.size(0)).uniform_(*self.opt['jpeg_range2'])
out = torch.clamp(out, 0, 1)
out = self.jpeger(out, quality=jpeg_p)
# resize back + the final sinc filter
mode = random.choice(['area', 'bilinear', 'bicubic'])
out = F.interpolate(out, size=(ori_h // self.opt['scale'], ori_w // self.opt['scale']), mode=mode)
out = filter2D(out, self.sinc_kernel)
# clamp and round
self.lq = torch.clamp((out * 255.0).round(), 0, 255) / 255.
# random crop
gt_size = self.opt['gt_size']
self.gt, self.lq = paired_random_crop(self.gt, self.lq, gt_size, self.opt['scale'])
# training pair pool
self._dequeue_and_enqueue()
self.lq = self.lq.contiguous() # for the warning: grad and param do not obey the gradient layout contract
else:
# for paired training or validation
self.lq = data['lq'].to(self.device)
if 'gt' in data:
self.gt = data['gt'].to(self.device)
self.gt_usm = self.usm_sharpener(self.gt)
# self.gt = self.usm_sharpener(self.gt)
def nondist_validation(self, dataloader, current_iter, tb_logger, save_img):
# do not use the synthetic process during validation
self.is_train = False
super(KDSRNetTAModel, self).nondist_validation(dataloader, current_iter, tb_logger, save_img)
self.is_train = True
def test(self):
if hasattr(self, 'net_g_ema'):
self.net_g_ema.eval()
with torch.no_grad():
self.output = self.net_g_ema(self.lq, self.gt)
else:
self.net_g.eval()
with torch.no_grad():
# self.output = self.net_g(self.lq, self.gt)
self.output = self.net_g(self.lq, self.gt)
self.net_g.train()
def optimize_parameters(self, current_iter):
self.optimizer_g.zero_grad()
self.output, _ = self.net_g(self.lq, self.gt)
l_total = 0
loss_dict = OrderedDict()
# pixel loss
if self.cri_pix:
l_pix = self.cri_pix(self.output, self.gt)
l_total += l_pix
loss_dict['l_pix'] = l_pix
# perceptual loss
if self.cri_perceptual:
l_percep, l_style = self.cri_perceptual(self.output, self.gt)
if l_percep is not None:
l_total += l_percep
loss_dict['l_percep'] = l_percep
if l_style is not None:
l_total += l_style
loss_dict['l_style'] = l_style
l_total.backward()
self.optimizer_g.step()
self.log_dict = self.reduce_loss_dict(loss_dict)
if self.ema_decay > 0:
self.model_ema(decay=self.ema_decay)
================================================
FILE: KDSR-GAN/kdsrgan/test.py
================================================
# flake8: noqa
import os.path as osp
from basicsr.test import test_pipeline
import kdsrgan.archs
import kdsrgan.data
import kdsrgan.models
if __name__ == '__main__':
root_path = osp.abspath(osp.join(__file__, osp.pardir, osp.pardir))
test_pipeline(root_path)
================================================
FILE: KDSR-GAN/kdsrgan/train.py
================================================
# flake8: noqa
import os.path as osp
from basicsr.train import train_pipeline
import kdsrgan.archs
import kdsrgan.data
import kdsrgan.models
import kdsrgan.losses
import warnings
warnings.filterwarnings("ignore")
if __name__ == '__main__':
root_path = osp.abspath(osp.join(__file__, osp.pardir, osp.pardir))
train_pipeline(root_path)
================================================
FILE: KDSR-GAN/kdsrgan/utils.py
================================================
import cv2
import math
import numpy as np
import os
import queue
import threading
import torch
from basicsr.utils.download_util import load_file_from_url
from torch.nn import functional as F
ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
class RealESRGANer():
"""A helper class for upsampling images with RealESRGAN.
Args:
scale (int): Upsampling scale factor used in the networks. It is usually 2 or 4.
model_path (str): The path to the pretrained model. It can be urls (will first download it automatically).
model (nn.Module): The defined network. Default: None.
tile (int): As too large images result in the out of GPU memory issue, so this tile option will first crop
input images into tiles, and then process each of them. Finally, they will be merged into one image.
0 denotes for do not use tile. Default: 0.
tile_pad (int): The pad size for each tile, to remove border artifacts. Default: 10.
pre_pad (int): Pad the input images to avoid border artifacts. Default: 10.
half (float): Whether to use half precision during inference. Default: False.
"""
def __init__(self,
scale,
model_path,
model=None,
tile=0,
tile_pad=10,
pre_pad=10,
half=False,
device=None,
gpu_id=None):
self.scale = scale
self.tile_size = tile
self.tile_pad = tile_pad
self.pre_pad = pre_pad
self.mod_scale = None
self.half = half
# initialize model
if gpu_id:
self.device = torch.device(
f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu') if device is None else device
else:
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device
# if the model_path starts with https, it will first download models to the folder: realesrgan/weights
if model_path.startswith('https://'):
model_path = load_file_from_url(
url=model_path, model_dir=os.path.join(ROOT_DIR, 'realesrgan/weights'), progress=True, file_name=None)
loadnet = torch.load(model_path, map_location=torch.device('cpu'))
# prefer to use params_ema
if 'params_ema' in loadnet:
keyname = 'params_ema'
else:
keyname = 'params'
model.load_state_dict(loadnet[keyname], strict=True)
model.eval()
self.model = model.to(self.device)
if self.half:
self.model = self.model.half()
def pre_process(self, img):
"""Pre-process, such as pre-pad and mod pad, so that the images can be divisible
"""
img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()
self.img = img.unsqueeze(0).to(self.device)
if self.half:
self.img = self.img.half()
# pre_pad
if self.pre_pad != 0:
self.img = F.pad(self.img, (0, self.pre_pad, 0, self.pre_pad), 'reflect')
# mod pad for divisible borders
if self.scale == 2:
self.mod_scale = 2
elif self.scale == 1:
self.mod_scale = 4
if self.mod_scale is not None:
self.mod_pad_h, self.mod_pad_w = 0, 0
_, _, h, w = self.img.size()
if (h % self.mod_scale != 0):
self.mod_pad_h = (self.mod_scale - h % self.mod_scale)
if (w % self.mod_scale != 0):
self.mod_pad_w = (self.mod_scale - w % self.mod_scale)
self.img = F.pad(self.img, (0, self.mod_pad_w, 0, self.mod_pad_h), 'reflect')
def process(self):
# model inference
self.output = self.model(self.img)
def tile_process(self):
"""It will first crop input images to tiles, and then process each tile.
Finally, all the processed tiles are merged into one images.
Modified from: https://github.com/ata4/esrgan-launcher
"""
batch, channel, height, width = self.img.shape
output_height = height * self.scale
output_width = width * self.scale
output_shape = (batch, channel, output_height, output_width)
# start with black image
self.output = self.img.new_zeros(output_shape)
tiles_x = math.ceil(width / self.tile_size)
tiles_y = math.ceil(height / self.tile_size)
# loop over all tiles
for y in range(tiles_y):
for x in range(tiles_x):
# extract tile from input image
ofs_x = x * self.tile_size
ofs_y = y * self.tile_size
# input tile area on total image
input_start_x = ofs_x
input_end_x = min(ofs_x + self.tile_size, width)
input_start_y = ofs_y
input_end_y = min(ofs_y + self.tile_size, height)
# input tile area on total image with padding
input_start_x_pad = max(input_start_x - self.tile_pad, 0)
input_end_x_pad = min(input_end_x + self.tile_pad, width)
input_start_y_pad = max(input_start_y - self.tile_pad, 0)
input_end_y_pad = min(input_end_y + self.tile_pad, height)
# input tile dimensions
input_tile_width = input_end_x - input_start_x
input_tile_height = input_end_y - input_start_y
tile_idx = y * tiles_x + x + 1
input_tile = self.img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad]
# upscale tile
try:
with torch.no_grad():
output_tile = self.model(input_tile)
except RuntimeError as error:
print('Error', error)
print(f'\tTile {tile_idx}/{tiles_x * tiles_y}')
# output tile area on total image
output_start_x = input_start_x * self.scale
output_end_x = input_end_x * self.scale
output_start_y = input_start_y * self.scale
output_end_y = input_end_y * self.scale
# output tile area without padding
output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale
output_end_x_tile = output_start_x_tile + input_tile_width * self.scale
output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale
output_end_y_tile = output_start_y_tile + input_tile_height * self.scale
# put tile into output image
self.output[:, :, output_start_y:output_end_y,
output_start_x:output_end_x] = output_tile[:, :, output_start_y_tile:output_end_y_tile,
output_start_x_tile:output_end_x_tile]
def post_process(self):
# remove extra pad
if self.mod_scale is not None:
_, _, h, w = self.output.size()
self.output = self.output[:, :, 0:h - self.mod_pad_h * self.scale, 0:w - self.mod_pad_w * self.scale]
# remove prepad
if self.pre_pad != 0:
_, _, h, w = self.output.size()
self.output = self.output[:, :, 0:h - self.pre_pad * self.scale, 0:w - self.pre_pad * self.scale]
return self.output
@torch.no_grad()
def enhance(self, img, outscale=None, alpha_upsampler='realesrgan'):
h_input, w_input = img.shape[0:2]
# img: numpy
img = img.astype(np.float32)
if np.max(img) > 256: # 16-bit image
max_range = 65535
print('\tInput is a 16-bit image')
else:
max_range = 255
img = img / max_range
if len(img.shape) == 2: # gray image
img_mode = 'L'
img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
elif img.shape[2] == 4: # RGBA image with alpha channel
img_mode = 'RGBA'
alpha = img[:, :, 3]
img = img[:, :, 0:3]
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
if alpha_upsampler == 'realesrgan':
alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2RGB)
else:
img_mode = 'RGB'
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# ------------------- process image (without the alpha channel) ------------------- #
self.pre_process(img)
if self.tile_size > 0:
self.tile_process()
else:
self.process()
output_img = self.post_process()
output_img = output_img.data.squeeze().float().cpu().clamp_(0, 1).numpy()
output_img = np.transpose(output_img[[2, 1, 0], :, :], (1, 2, 0))
if img_mode == 'L':
output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2GRAY)
# ------------------- process the alpha channel if necessary ------------------- #
if img_mode == 'RGBA':
if alpha_upsampler == 'realesrgan':
self.pre_process(alpha)
if self.tile_size > 0:
self.tile_process()
else:
self.process()
output_alpha = self.post_process()
output_alpha = output_alpha.data.squeeze().float().cpu().clamp_(0, 1).numpy()
output_alpha = np.transpose(output_alpha[[2, 1, 0], :, :], (1, 2, 0))
output_alpha = cv2.cvtColor(output_alpha, cv2.COLOR_BGR2GRAY)
else: # use the cv2 resize for alpha channel
h, w = alpha.shape[0:2]
output_alpha = cv2.resize(alpha, (w * self.scale, h * self.scale), interpolation=cv2.INTER_LINEAR)
# merge the alpha channel
output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2BGRA)
output_img[:, :, 3] = output_alpha
# ------------------------------ return ------------------------------ #
if max_range == 65535: # 16-bit image
output = (output_img * 65535.0).round().astype(np.uint16)
else:
output = (output_img * 255.0).round().astype(np.uint8)
if outscale is not None and outscale != float(self.scale):
output = cv2.resize(
output, (
int(w_input * outscale),
int(h_input * outscale),
), interpolation=cv2.INTER_LANCZOS4)
return output, img_mode
class PrefetchReader(threading.Thread):
"""Prefetch images.
Args:
img_list (list[str]): A image list of image paths to be read.
num_prefetch_queue (int): Number of prefetch queue.
"""
def __init__(self, img_list, num_prefetch_queue):
super().__init__()
self.que = queue.Queue(num_prefetch_queue)
self.img_list = img_list
def run(self):
for img_path in self.img_list:
img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED)
self.que.put(img)
self.que.put(None)
def __next__(self):
next_item = self.que.get()
if next_item is None:
raise StopIteration
return next_item
def __iter__(self):
return self
class IOConsumer(threading.Thread):
def __init__(self, opt, que, qid):
super().__init__()
self._queue = que
self.qid = qid
self.opt = opt
def run(self):
while True:
msg = self._queue.get()
if isinstance(msg, str) and msg == 'quit':
break
output = msg['output']
save_path = msg['save_path']
cv2.imwrite(save_path, output)
print(f'IO worker {self.qid} is done.')
================================================
FILE: KDSR-GAN/kdsrgan/weights/README.md
================================================
# Weights
Put the downloaded weights to this folder.
================================================
FILE: KDSR-GAN/options/test_kdsrgan_x4ST.yml
================================================
# general settings
name: test_KDSRGANx4STplus_400k_B12G4
model_type: KDSRGANSTModel
scale: 4
num_gpu: auto # auto: can infer from your visible devices automatically. official: 4 GPUs
manual_seed: 0
# ----------------- options for synthesizing training data in RealESRGANModel ----------------- #
# dataset and data loader settings
datasets:
# Uncomment these for validation
val_1:
name: NTIRE2020-Track1
type: PairedImageDataset
dataroot_gt: /root/dataset/NTIRE2020-Track1/track1-valid-gt
dataroot_lq: /root/dataset/NTIRE2020-Track1/track1-valid-input
io_backend:
type: disk
# network structures
network_g:
type: BlindSR_ST
n_feats: 128
n_encoder_res: 6
scale: 4
n_sr_blocks: 42
# network structures
network_TA:
type: BlindSR_TA
n_feats: 128
n_encoder_res: 6
scale: 4
n_sr_blocks: 42
network_d:
type: UNetDiscriminatorSN
num_in_ch: 3
num_feat: 64
skip_connection: True
# path
path:
# use the pre-trained Real-ESRNet model
pretrain_network_TA: experiments/KDSRT-rec.pth
pretrain_network_g: experiments/KDSRS-GAN.pth
# pretrain_network_g: experiments/KDSRS-GANV2.pth
param_key_g: params_ema
strict_load_g: False
ignore_resume_networks: network_TA
val:
save_img: True
suffix: ~ # add suffix to saved images, if None, use exp name
metrics:
psnr: # metric name
type: calculate_psnr
crop_border: 4
test_y_channel: true
================================================
FILE: KDSR-GAN/options/test_kdsrnet_x4TA.yml
================================================
# general settings
name: train_KDSRNetTAx4plus_1000k_B12G4
model_type: KDSRNetTAModel
scale: 4
num_gpu: auto # auto: can infer from your visible devices automatically. official: 4 GPUs
manual_seed: 0
# ----------------- options for synthesizing training data in RealESRNetModel ----------------- #
gt_usm: True # USM the ground-truth
# dataset and data loader settings
datasets:
# Uncomment these for validation
test_1:
name: NTIRE2020-Track1
type: PairedImageDataset
dataroot_gt: /root/dataset/NTIRE2020-Track1/track1-valid-gt
dataroot_lq: /root/dataset/NTIRE2020-Track1/track1-valid-input
io_backend:
type: disk
# network structures
network_g:
type: BlindSR_TA
n_feats: 128
n_encoder_res: 6
scale: 4
n_sr_blocks: 42
# path
path:
pretrain_network_g: ./experiments/KDSRT-rec.pth
param_key_g: params_ema
strict_load_g: true
# Uncomment these for validation
# validation settings
val:
save_img: True
suffix: ~ # add suffix to saved images, if None, use exp name
metrics:
psnr: # metric name
type: calculate_psnr
crop_border: 4
test_y_channel: true
ssim:
type: calculate_ssim
crop_border: 4
test_y_channel: true
================================================
FILE: KDSR-GAN/options/train_kdsrgan_x4ST.yml
================================================
# general settings
name: train_KDSRGANx4STplus_400k_B12G4
model_type: KDSRGANSTModel
scale: 4
num_gpu: auto # auto: can infer from your visible devices automatically. official: 4 GPUs
manual_seed: 0
# ----------------- options for synthesizing training data in RealESRGANModel ----------------- #
l1_gt_usm: False
percep_gt_usm: False
gan_gt_usm: False
# the first degradation process
resize_prob: [0.2, 0.7, 0.1] # up, down, keep
resize_range: [0.15, 1.5]
gaussian_noise_prob: 0.5
noise_range: [1, 30]
poisson_scale_range: [0.05, 3]
gray_noise_prob: 0.4
jpeg_range: [30, 95]
# the second degradation process
second_blur_prob: 0.8
resize_prob2: [0.3, 0.4, 0.3] # up, down, keep
resize_range2: [0.3, 1.2]
gaussian_noise_prob2: 0.5
noise_range2: [1, 25]
poisson_scale_range2: [0.05, 2.5]
gray_noise_prob2: 0.4
jpeg_range2: [30, 95]
gt_size: 256
queue_size: 180
# dataset and data loader settings
datasets:
train:
name: DF2K+OST
type: RealESRGANDataset
dataroot_gt: /root/dataset
meta_info: datasets/meta_info/meta_info_DF2Kmultiscale+OST_sub.txt
io_backend:
type: disk
blur_kernel_size: 21
kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
sinc_prob: 0.1
blur_sigma: [0.2, 3]
betag_range: [0.5, 4]
betap_range: [1, 2]
blur_kernel_size2: 21
kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']
kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]
sinc_prob2: 0.1
blur_sigma2: [0.2, 1.5]
betag_range
gitextract_r69x67a9/ ├── KDSR-GAN/ │ ├── .gitignore │ ├── .pre-commit-config.yaml │ ├── CODE_OF_CONDUCT.md │ ├── Metric/ │ │ ├── DISTS/ │ │ │ ├── DISTS_pytorch/ │ │ │ │ ├── DISTS_pt.py │ │ │ │ └── weights.pt │ │ │ ├── DISTS_tensorflow/ │ │ │ │ └── DISTS_tf.py │ │ │ ├── LICENSE │ │ │ ├── requirements.txt │ │ │ └── weights/ │ │ │ └── alpha_beta.mat │ │ ├── LPIPS.py │ │ ├── PSNR.py │ │ ├── dists.py │ │ └── pip.sh │ ├── README.md │ ├── VERSION │ ├── docs/ │ │ ├── CONTRIBUTING.md │ │ ├── FAQ.md │ │ ├── Training.md │ │ ├── Training_CN.md │ │ ├── anime_comparisons.md │ │ ├── anime_comparisons_CN.md │ │ ├── anime_model.md │ │ ├── anime_video_model.md │ │ ├── feedback.md │ │ ├── model_zoo.md │ │ └── ncnn_conversion.md │ ├── kdsrgan/ │ │ ├── __init__.py │ │ ├── archs/ │ │ │ ├── ST_arch.py │ │ │ ├── TA_arch.py │ │ │ ├── __init__.py │ │ │ ├── common.py │ │ │ ├── discriminator_arch.py │ │ │ └── srvgg_arch.py │ │ ├── data/ │ │ │ ├── __init__.py │ │ │ ├── kdsrgan_dataset.py │ │ │ └── kdsrgan_paired_dataset.py │ │ ├── losses/ │ │ │ ├── __init__.py │ │ │ └── my_loss.py │ │ ├── models/ │ │ │ ├── __init__.py │ │ │ ├── kdsrgan_ST_model.py │ │ │ ├── kdsrgan_TA_model.py │ │ │ ├── kdsrnet_ST_model.py │ │ │ └── kdsrnet_TA_model.py │ │ ├── test.py │ │ ├── train.py │ │ ├── utils.py │ │ └── weights/ │ │ └── README.md │ ├── options/ │ │ ├── test_kdsrgan_x4ST.yml │ │ ├── test_kdsrnet_x4TA.yml │ │ ├── train_kdsrgan_x4ST.yml │ │ ├── train_kdsrgan_x4STV2.yml │ │ ├── train_kdsrnet_x4ST.yml │ │ └── train_kdsrnet_x4TA.yml │ ├── pip.sh │ ├── requirements.txt │ ├── scripts/ │ │ ├── extract_subimages.py │ │ ├── extract_subimages_DF2K.py │ │ ├── generate_meta_info.py │ │ ├── generate_meta_info_DF2K.py │ │ ├── generate_meta_info_OST.py │ │ ├── generate_meta_info_pairdata.py │ │ ├── generate_multiscale_DF2K.py │ │ └── pytorch2onnx.py │ ├── setup.cfg │ ├── setup.py │ ├── test.sh │ ├── tests/ │ │ ├── data/ │ │ │ ├── gt.lmdb/ │ │ │ │ ├── data.mdb │ │ │ │ ├── lock.mdb │ │ │ │ └── meta_info.txt │ │ │ ├── lq.lmdb/ │ │ │ │ ├── data.mdb │ │ │ │ ├── lock.mdb │ │ │ │ └── meta_info.txt │ │ │ ├── meta_info_gt.txt │ │ │ ├── meta_info_pair.txt │ │ │ ├── test_realesrgan_dataset.yml │ │ │ ├── test_realesrgan_model.yml │ │ │ ├── test_realesrgan_paired_dataset.yml │ │ │ └── test_realesrnet_model.yml │ │ ├── test_dataset.py │ │ ├── test_discriminator_arch.py │ │ ├── test_model.py │ │ └── test_utils.py │ ├── train_KDSRGAN_S.sh │ ├── train_KDSRNet_S.sh │ └── train_KDSRNet_T.sh ├── KDSR-classic/ │ ├── README.md │ ├── cal_parms.py │ ├── data/ │ │ ├── __init__.py │ │ ├── benchmark.py │ │ ├── common.py │ │ ├── common2.py │ │ ├── df2k.py │ │ ├── multiscalesrdata.py │ │ └── multiscalesrdata2.py │ ├── dataloader.py │ ├── loss/ │ │ ├── __init__.py │ │ ├── adversarial.py │ │ ├── discriminator.py │ │ └── vgg.py │ ├── main_anisonoise_KDSRsMx4_stage3.sh │ ├── main_anisonoise_KDSRsMx4_stage4.sh │ ├── main_anisonoise_stage3.py │ ├── main_anisonoise_stage4.py │ ├── main_iso_KDSRsLx4_stage3.sh │ ├── main_iso_KDSRsLx4_stage4.sh │ ├── main_iso_KDSRsMx4_stage3.sh │ ├── main_iso_KDSRsMx4_stage4.sh │ ├── main_iso_stage3.py │ ├── main_iso_stage4.py │ ├── model/ │ │ ├── __init__.py │ │ ├── blindsr.py │ │ └── common.py │ ├── model_ST/ │ │ ├── __init__.py │ │ ├── blindsr.py │ │ └── common.py │ ├── model_TA/ │ │ ├── __init__.py │ │ ├── blindsr.py │ │ └── common.py │ ├── option.py │ ├── template.py │ ├── test.py │ ├── test_anisonoise_KDSRsMx4.sh │ ├── test_iso_KDSRsLx4.sh │ ├── test_iso_KDSRsMx4.sh │ ├── trainer_anisonoise_stage3.py │ ├── trainer_anisonoise_stage4.py │ ├── trainer_iso_stage3.py │ ├── trainer_iso_stage4.py │ ├── utility.py │ ├── utility2.py │ ├── utility3.py │ └── utils/ │ ├── __init__.py │ ├── util.py │ └── util2.py └── README.md
SYMBOL INDEX (447 symbols across 65 files)
FILE: KDSR-GAN/Metric/DISTS/DISTS_pytorch/DISTS_pt.py
class L2pooling (line 11) | class L2pooling(nn.Module):
method __init__ (line 12) | def __init__(self, filter_size=5, stride=2, channels=None, pad_off=0):
method forward (line 22) | def forward(self, input):
class DISTS (line 27) | class DISTS(torch.nn.Module):
method __init__ (line 28) | def __init__(self, load_weights=True):
method forward_once (line 68) | def forward_once(self, x):
method forward (line 82) | def forward(self, x, y, require_grad=False, batch_average=False):
function prepare_image (line 115) | def prepare_image(image, resize=True):
FILE: KDSR-GAN/Metric/DISTS/DISTS_tensorflow/DISTS_tf.py
class DISTS (line 13) | class DISTS():
method __init__ (line 14) | def __init__(self):
method get_features (line 25) | def get_features(self, img):
method conv_layer (line 53) | def conv_layer(self, input, name):
method pool_layer (line 61) | def pool_layer(self, input, name):
method get_conv_filter (line 68) | def get_conv_filter(self, name):
method get_bias (line 71) | def get_bias(self, name):
method get_score (line 74) | def get_score(self, img1, img2):
FILE: KDSR-GAN/Metric/LPIPS.py
function main (line 11) | def main():
FILE: KDSR-GAN/Metric/PSNR.py
function main (line 12) | def main():
FILE: KDSR-GAN/Metric/dists.py
class L2pooling (line 13) | class L2pooling(nn.Module):
method __init__ (line 14) | def __init__(self, filter_size=5, stride=2, channels=None, pad_off=0):
method forward (line 24) | def forward(self, input):
class DISTS (line 29) | class DISTS(torch.nn.Module):
method __init__ (line 30) | def __init__(self, load_weights=True):
method forward_once (line 70) | def forward_once(self, x):
method forward (line 84) | def forward(self, x, y, require_grad=False, batch_average=False):
function prepare_image (line 117) | def prepare_image(image, resize=True):
FILE: KDSR-GAN/kdsrgan/archs/ST_arch.py
class IDR_DDC (line 8) | class IDR_DDC(nn.Module):
method __init__ (line 9) | def __init__(self, channels_in, channels_out, kernel_size, reduction):
method forward (line 21) | def forward(self, x):
class IDR_DCRB (line 38) | class IDR_DCRB(nn.Module):
method __init__ (line 39) | def __init__(self, conv, n_feat, kernel_size, reduction):
method forward (line 46) | def forward(self, x):
class DAG (line 59) | class DAG(nn.Module):
method __init__ (line 60) | def __init__(self, conv, n_feat, kernel_size, reduction, n_blocks):
method forward (line 71) | def forward(self, x):
class KDSR (line 83) | class KDSR(nn.Module):
method __init__ (line 84) | def __init__(self, n_feats=128, scale=4,n_sr_blocks = 42, rgb_range=1,...
method forward (line 111) | def forward(self, x, k_v):
class KD_IDE (line 135) | class KD_IDE(nn.Module):
method __init__ (line 136) | def __init__(self,n_feats = 128, n_encoder_res = 6):
method forward (line 169) | def forward(self, x):
class BlindSR_ST (line 178) | class BlindSR_ST(nn.Module):
method __init__ (line 179) | def __init__(self, n_feats=128, n_encoder_res=6, scale=4,n_sr_blocks=4...
method forward (line 190) | def forward(self, x):
FILE: KDSR-GAN/kdsrgan/archs/TA_arch.py
class IDR_DDC (line 7) | class IDR_DDC(nn.Module):
method __init__ (line 8) | def __init__(self, channels_in, channels_out, kernel_size, reduction):
method forward (line 20) | def forward(self, x):
class IDR_DCRB (line 37) | class IDR_DCRB(nn.Module):
method __init__ (line 38) | def __init__(self, conv, n_feat, kernel_size, reduction):
method forward (line 45) | def forward(self, x):
class DAG (line 58) | class DAG(nn.Module):
method __init__ (line 59) | def __init__(self, conv, n_feat, kernel_size, reduction, n_blocks):
method forward (line 70) | def forward(self, x):
class KDSR (line 84) | class KDSR(nn.Module):
method __init__ (line 85) | def __init__(self, n_feats=128, scale=4,n_sr_blocks = 42, rgb_range=1,...
method forward (line 112) | def forward(self, x, k_v):
class KD_IDE (line 136) | class KD_IDE(nn.Module):
method __init__ (line 137) | def __init__(self,n_feats = 128, n_encoder_res = 6):
method forward (line 170) | def forward(self, x):
class BlindSR_TA (line 179) | class BlindSR_TA(nn.Module):
method __init__ (line 180) | def __init__(self, n_feats=128, n_encoder_res=6, scale=4,n_sr_blocks=4...
method forward (line 191) | def forward(self, x, gt):
FILE: KDSR-GAN/kdsrgan/archs/common.py
function default_conv (line 7) | def default_conv(in_channels, out_channels, kernel_size, bias=True):
class ResBlock (line 10) | class ResBlock(nn.Module):
method __init__ (line 11) | def __init__(
method forward (line 27) | def forward(self, x):
class MeanShift (line 33) | class MeanShift(nn.Conv2d):
method __init__ (line 34) | def __init__(self, rgb_range, rgb_mean, rgb_std, sign=-1):
class Upsampler (line 45) | class Upsampler(nn.Sequential):
method __init__ (line 46) | def __init__(self, conv, scale, n_feat, act=False, bias=True):
FILE: KDSR-GAN/kdsrgan/archs/discriminator_arch.py
class UNetDiscriminatorSN (line 8) | class UNetDiscriminatorSN(nn.Module):
method __init__ (line 19) | def __init__(self, num_in_ch, num_feat=64, skip_connection=True):
method forward (line 38) | def forward(self, x):
FILE: KDSR-GAN/kdsrgan/archs/srvgg_arch.py
class SRVGGNetCompact (line 7) | class SRVGGNetCompact(nn.Module):
method __init__ (line 22) | def __init__(self, num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=16...
method forward (line 60) | def forward(self, x):
FILE: KDSR-GAN/kdsrgan/data/kdsrgan_dataset.py
class KDSRGANDataset (line 17) | class KDSRGANDataset(data.Dataset):
method __init__ (line 35) | def __init__(self, opt):
method __getitem__ (line 83) | def __getitem__(self, index):
method __len__ (line 191) | def __len__(self):
FILE: KDSR-GAN/kdsrgan/data/kdsrgan_paired_dataset.py
class KDSRGANPairedDataset (line 11) | class KDSRGANPairedDataset(data.Dataset):
method __init__ (line 41) | def __init__(self, opt):
method __getitem__ (line 75) | def __getitem__(self, index):
method __len__ (line 107) | def __len__(self):
FILE: KDSR-GAN/kdsrgan/losses/my_loss.py
class KDLoss (line 8) | class KDLoss(nn.Module):
method __init__ (line 14) | def __init__(self, loss_weight=1.0, temperature = 0.15):
method forward (line 20) | def forward(self, T_fea, S_fea):
FILE: KDSR-GAN/kdsrgan/models/kdsrgan_ST_model.py
class KDSRGANSTModel (line 18) | class KDSRGANSTModel(SRGANModel):
method __init__ (line 26) | def __init__(self, opt):
method init_training_settings (line 58) | def init_training_settings(self):
method _dequeue_and_enqueue (line 69) | def _dequeue_and_enqueue(self):
method feed_data (line 106) | def feed_data(self, data):
method nondist_validation (line 224) | def nondist_validation(self, dataloader, current_iter, tb_logger, save...
method test (line 230) | def test(self):
method optimize_parameters (line 241) | def optimize_parameters(self, current_iter):
FILE: KDSR-GAN/kdsrgan/models/kdsrgan_TA_model.py
class KDSRGANTAModel (line 15) | class KDSRGANTAModel(SRGANModel):
method __init__ (line 23) | def __init__(self, opt):
method _dequeue_and_enqueue (line 30) | def _dequeue_and_enqueue(self):
method feed_data (line 67) | def feed_data(self, data):
method nondist_validation (line 185) | def nondist_validation(self, dataloader, current_iter, tb_logger, save...
method optimize_parameters (line 191) | def optimize_parameters(self, current_iter):
FILE: KDSR-GAN/kdsrgan/models/kdsrnet_ST_model.py
class KDSRNetSTModel (line 19) | class KDSRNetSTModel(SRModel):
method __init__ (line 27) | def __init__(self, opt):
method setup_optimizers (line 59) | def setup_optimizers(self):
method setup_schedulers (line 73) | def setup_schedulers(self):
method init_training_settings (line 87) | def init_training_settings(self):
method _dequeue_and_enqueue (line 131) | def _dequeue_and_enqueue(self):
method feed_data (line 168) | def feed_data(self, data):
method nondist_validation (line 285) | def nondist_validation(self, dataloader, current_iter, tb_logger, save...
method test (line 291) | def test(self):
method optimize_parameters (line 303) | def optimize_parameters(self, current_iter):
FILE: KDSR-GAN/kdsrgan/models/kdsrnet_TA_model.py
class KDSRNetTAModel (line 15) | class KDSRNetTAModel(SRModel):
method __init__ (line 23) | def __init__(self, opt):
method _dequeue_and_enqueue (line 32) | def _dequeue_and_enqueue(self):
method feed_data (line 69) | def feed_data(self, data):
method nondist_validation (line 187) | def nondist_validation(self, dataloader, current_iter, tb_logger, save...
method test (line 193) | def test(self):
method optimize_parameters (line 205) | def optimize_parameters(self, current_iter):
FILE: KDSR-GAN/kdsrgan/utils.py
class RealESRGANer (line 14) | class RealESRGANer():
method __init__ (line 29) | def __init__(self,
method pre_process (line 68) | def pre_process(self, img):
method process (line 93) | def process(self):
method tile_process (line 97) | def tile_process(self):
method post_process (line 162) | def post_process(self):
method enhance (line 174) | def enhance(self, img, outscale=None, alpha_upsampler='realesrgan'):
class PrefetchReader (line 246) | class PrefetchReader(threading.Thread):
method __init__ (line 254) | def __init__(self, img_list, num_prefetch_queue):
method run (line 259) | def run(self):
method __next__ (line 266) | def __next__(self):
method __iter__ (line 272) | def __iter__(self):
class IOConsumer (line 276) | class IOConsumer(threading.Thread):
method __init__ (line 278) | def __init__(self, opt, que, qid):
method run (line 284) | def run(self):
FILE: KDSR-GAN/scripts/extract_subimages.py
function main (line 12) | def main(args):
function extract_subimages (line 43) | def extract_subimages(opt):
function worker (line 74) | def worker(path, opt):
FILE: KDSR-GAN/scripts/extract_subimages_DF2K.py
function main (line 12) | def main(args):
function extract_subimages (line 43) | def extract_subimages(opt):
function worker (line 74) | def worker(path, opt):
FILE: KDSR-GAN/scripts/generate_meta_info.py
function main (line 7) | def main(args):
FILE: KDSR-GAN/scripts/generate_meta_info_DF2K.py
function main (line 7) | def main(args):
FILE: KDSR-GAN/scripts/generate_meta_info_OST.py
function main (line 7) | def main(args):
FILE: KDSR-GAN/scripts/generate_meta_info_pairdata.py
function main (line 6) | def main(args):
FILE: KDSR-GAN/scripts/generate_multiscale_DF2K.py
function main (line 7) | def main(args):
FILE: KDSR-GAN/scripts/pytorch2onnx.py
function main (line 7) | def main(args):
FILE: KDSR-GAN/setup.py
function readme (line 12) | def readme():
function get_git_hash (line 18) | def get_git_hash():
function get_hash (line 43) | def get_hash():
function write_version_py (line 52) | def write_version_py():
function get_version (line 69) | def get_version():
function get_requirements (line 75) | def get_requirements(filename='requirements.txt'):
FILE: KDSR-GAN/tests/test_dataset.py
function test_realesrgan_dataset (line 8) | def test_realesrgan_dataset():
function test_realesrgan_paired_dataset (line 82) | def test_realesrgan_paired_dataset():
FILE: KDSR-GAN/tests/test_discriminator_arch.py
function test_unetdiscriminatorsn (line 6) | def test_unetdiscriminatorsn():
FILE: KDSR-GAN/tests/test_model.py
function test_realesrnet_model (line 12) | def test_realesrnet_model():
function test_realesrgan_model (line 64) | def test_realesrgan_model():
FILE: KDSR-GAN/tests/test_utils.py
function test_realesrganer (line 7) | def test_realesrganer():
FILE: KDSR-classic/cal_parms.py
function count_param (line 1) | def count_param(model):
FILE: KDSR-classic/data/__init__.py
class MyConcatDataset (line 7) | class MyConcatDataset(ConcatDataset):
method __init__ (line 8) | def __init__(self, datasets):
method set_scale (line 12) | def set_scale(self, idx_scale):
class Data (line 16) | class Data:
method __init__ (line 17) | def __init__(self, args):
FILE: KDSR-classic/data/benchmark.py
class Benchmark (line 6) | class Benchmark(srdata.SRData):
method __init__ (line 7) | def __init__(self, args, name='', train=True):
method _set_filesystem (line 12) | def _set_filesystem(self, dir_data):
FILE: KDSR-classic/data/common.py
function get_patch (line 9) | def get_patch(hr, patch_size=48, scale=2, multi=False, input_large=False):
function set_channel (line 22) | def set_channel(hr, n_channels=3):
function np2Tensor (line 38) | def np2Tensor(hr, rgb_range=255):
function augment (line 49) | def augment(hr, hflip=True, rot=True):
FILE: KDSR-classic/data/common2.py
function get_patch (line 9) | def get_patch(*args, patch_size=96, scale=2, multi=False, input_large=Fa...
function set_channel (line 36) | def set_channel(*args, n_channels=3):
function np2Tensor (line 52) | def np2Tensor(*args, rgb_range=255):
function augment (line 63) | def augment(*args, hflip=True, rot=True):
FILE: KDSR-classic/data/df2k.py
class DF2K (line 5) | class DF2K(multiscalesrdata.SRData):
method __init__ (line 6) | def __init__(self, args, name='DF2K', train=True, benchmark=False):
method _scan (line 19) | def _scan(self):
method _set_filesystem (line 25) | def _set_filesystem(self, dir_data):
FILE: KDSR-classic/data/multiscalesrdata.py
class SRData (line 14) | class SRData(data.Dataset):
method __init__ (line 15) | def __init__(self, args, name='', train=True, benchmark=False):
method _scan (line 69) | def _scan(self):
method _set_filesystem (line 85) | def _set_filesystem(self, dir_data):
method _check_and_load (line 92) | def _check_and_load(self, ext, img, f, verbose=True):
method __getitem__ (line 99) | def __getitem__(self, idx):
method __len__ (line 107) | def __len__(self):
method _get_index (line 113) | def _get_index(self, idx):
method _load_file (line 119) | def _load_file(self, idx):
method get_patch (line 132) | def get_patch(self, hr):
method set_scale (line 151) | def set_scale(self, idx_scale):
FILE: KDSR-classic/data/multiscalesrdata2.py
class SRData (line 14) | class SRData(data.Dataset):
method __init__ (line 15) | def __init__(self, args, name='', train=True, benchmark=False):
method _scan (line 69) | def _scan(self):
method _set_filesystem (line 86) | def _set_filesystem(self, dir_data):
method _check_and_load (line 93) | def _check_and_load(self, ext, img, f, verbose=True):
method __getitem__ (line 100) | def __getitem__(self, idx):
method __len__ (line 108) | def __len__(self):
method _get_index (line 114) | def _get_index(self, idx):
method _load_file (line 120) | def _load_file(self, idx):
method get_patch (line 137) | def get_patch(self, lr, hr):
method set_scale (line 154) | def set_scale(self, idx_scale):
FILE: KDSR-classic/dataloader.py
function _ms_loop (line 20) | def _ms_loop(dataset, index_queue, data_queue, collate_fn, scale, seed, ...
class _MSDataLoaderIter (line 46) | class _MSDataLoaderIter(_DataLoaderIter):
method __init__ (line 47) | def __init__(self, loader):
class MSDataLoader (line 117) | class MSDataLoader(DataLoader):
method __init__ (line 118) | def __init__(
method __iter__ (line 133) | def __iter__(self):
FILE: KDSR-classic/loss/__init__.py
class Loss (line 13) | class Loss(nn.modules.loss._Loss):
method __init__ (line 14) | def __init__(self, args, ckp):
method forward (line 70) | def forward(self, sr, hr):
method step (line 87) | def step(self):
method start_log (line 92) | def start_log(self):
method end_log (line 95) | def end_log(self, n_batches):
method display_loss (line 98) | def display_loss(self, batch):
method plot_loss (line 106) | def plot_loss(self, apath, epoch):
method get_loss_module (line 120) | def get_loss_module(self):
method save (line 126) | def save(self, apath):
method load (line 130) | def load(self, apath, cpu=False):
FILE: KDSR-classic/loss/adversarial.py
class Adversarial (line 11) | class Adversarial(nn.Module):
method __init__ (line 12) | def __init__(self, args, gan_type):
method forward (line 26) | def forward(self, fake, real):
method state_dict (line 78) | def state_dict(self, *args, **kwargs):
FILE: KDSR-classic/loss/discriminator.py
class Discriminator (line 5) | class Discriminator(nn.Module):
method __init__ (line 6) | def __init__(self, args, gan_type='GAN'):
method forward (line 40) | def forward(self, x):
FILE: KDSR-classic/loss/vgg.py
class VGG (line 9) | class VGG(nn.Module):
method __init__ (line 10) | def __init__(self, conv_index, rgb_range=1):
method forward (line 24) | def forward(self, sr, hr):
FILE: KDSR-classic/main_anisonoise_stage3.py
function count_param (line 10) | def count_param(model):
FILE: KDSR-classic/main_iso_stage3.py
function count_param (line 10) | def count_param(model):
FILE: KDSR-classic/model/__init__.py
class Model (line 8) | class Model(nn.Module):
method __init__ (line 9) | def __init__(self, args, ckp):
method forward (line 38) | def forward(self, lr,lr_bic):
method get_model (line 51) | def get_model(self):
method state_dict (line 57) | def state_dict(self, **kwargs):
method save (line 61) | def save(self, apath, epoch, is_best=False):
method load (line 79) | def load(self, apath, pre_train='.', resume=-1, cpu=False):
method forward_chop (line 104) | def forward_chop(self, x, shave=10, min_size=160000):
method forward_x8 (line 145) | def forward_x8(self, x, forward_function):
FILE: KDSR-classic/model/blindsr.py
function make_model (line 8) | def make_model(args):
class CZSR (line 12) | class CZSR(nn.Module):
method __init__ (line 13) | def __init__(self, args, conv=common.default_conv):
method forward (line 60) | def forward(self, lr,lr_bic):
FILE: KDSR-classic/model/common.py
function default_conv (line 7) | def default_conv(in_channels, out_channels, kernel_size, bias=True):
class ResBlock (line 10) | class ResBlock(nn.Module):
method __init__ (line 11) | def __init__(
method forward (line 27) | def forward(self, x):
class MeanShift (line 33) | class MeanShift(nn.Conv2d):
method __init__ (line 34) | def __init__(self, rgb_range, rgb_mean, rgb_std, sign=-1):
class Upsampler (line 45) | class Upsampler(nn.Sequential):
method __init__ (line 46) | def __init__(self, conv, scale, n_feat, act=False, bias=True):
FILE: KDSR-classic/model_ST/__init__.py
class Model (line 8) | class Model(nn.Module):
method __init__ (line 9) | def __init__(self, args, ckp):
method forward (line 38) | def forward(self, x):
method get_model (line 51) | def get_model(self):
method state_dict (line 57) | def state_dict(self, **kwargs):
method save (line 61) | def save(self, apath, epoch, is_best=False):
method load (line 79) | def load(self, apath, pre_train='.', resume=-1, cpu=False):
method forward_chop (line 104) | def forward_chop(self, x, shave=10, min_size=160000):
method forward_x8 (line 145) | def forward_x8(self, x, forward_function):
FILE: KDSR-classic/model_ST/blindsr.py
function make_model (line 7) | def make_model(args):
class IDR_DDC (line 10) | class IDR_DDC(nn.Module):
method __init__ (line 11) | def __init__(self, channels_in, channels_out, kernel_size, reduction):
method forward (line 23) | def forward(self, x):
class IDR_DCRB (line 40) | class IDR_DCRB(nn.Module):
method __init__ (line 41) | def __init__(self, conv, n_feat, kernel_size, reduction):
method forward (line 48) | def forward(self, x):
class DAG (line 61) | class DAG(nn.Module):
method __init__ (line 62) | def __init__(self, conv, n_feat, kernel_size, reduction, n_blocks):
method forward (line 73) | def forward(self, x):
class KDSR (line 85) | class KDSR(nn.Module):
method __init__ (line 86) | def __init__(self, args, conv=common.default_conv):
method forward (line 118) | def forward(self, x, k_v):
class KD_IDE (line 142) | class KD_IDE(nn.Module):
method __init__ (line 143) | def __init__(self,args):
method forward (line 180) | def forward(self, x):
class BlindSR (line 193) | class BlindSR(nn.Module):
method __init__ (line 194) | def __init__(self, args):
method forward (line 203) | def forward(self, x):
FILE: KDSR-classic/model_ST/common.py
function default_conv (line 7) | def default_conv(in_channels, out_channels, kernel_size, bias=True):
class MeanShift (line 11) | class MeanShift(nn.Conv2d):
method __init__ (line 12) | def __init__(self, rgb_range, rgb_mean, rgb_std, sign=-1):
class Upsampler (line 23) | class Upsampler(nn.Sequential):
method __init__ (line 24) | def __init__(self, conv, scale, n_feat, act=False, bias=True):
FILE: KDSR-classic/model_TA/__init__.py
class Model (line 8) | class Model(nn.Module):
method __init__ (line 9) | def __init__(self, args, ckp):
method forward (line 38) | def forward(self, x, deg_repre):
method get_model (line 51) | def get_model(self):
method state_dict (line 57) | def state_dict(self, **kwargs):
method save (line 61) | def save(self, apath, epoch, is_best=False):
method load (line 79) | def load(self, apath, pre_train='.', resume=-1, cpu=False):
method forward_chop (line 104) | def forward_chop(self, x, deg_repre, shave=10, min_size=160000):
method forward_x8 (line 156) | def forward_x8(self, x, forward_function):
FILE: KDSR-classic/model_TA/blindsr.py
function make_model (line 7) | def make_model(args):
class IDR_DDC (line 10) | class IDR_DDC(nn.Module):
method __init__ (line 11) | def __init__(self, channels_in, channels_out, kernel_size, reduction):
method forward (line 23) | def forward(self, x):
class IDR_DCRB (line 40) | class IDR_DCRB(nn.Module):
method __init__ (line 41) | def __init__(self, conv, n_feat, kernel_size, reduction):
method forward (line 48) | def forward(self, x):
class DAG (line 61) | class DAG(nn.Module):
method __init__ (line 62) | def __init__(self, conv, n_feat, kernel_size, reduction, n_blocks):
method forward (line 73) | def forward(self, x):
class KDSR (line 85) | class KDSR(nn.Module):
method __init__ (line 86) | def __init__(self, args, conv=common.default_conv):
method forward (line 118) | def forward(self, x, k_v):
class KD_IDE (line 142) | class KD_IDE(nn.Module):
method __init__ (line 143) | def __init__(self,args):
method forward (line 181) | def forward(self, x):
class BlindSR (line 190) | class BlindSR(nn.Module):
method __init__ (line 191) | def __init__(self, args):
method forward (line 200) | def forward(self, x, deg_repre):
FILE: KDSR-classic/model_TA/common.py
function default_conv (line 7) | def default_conv(in_channels, out_channels, kernel_size, bias=True):
class MeanShift (line 11) | class MeanShift(nn.Conv2d):
method __init__ (line 12) | def __init__(self, rgb_range, rgb_mean, rgb_std, sign=-1):
class Upsampler (line 23) | class Upsampler(nn.Sequential):
method __init__ (line 24) | def __init__(self, conv, scale, n_feat, act=False, bias=True):
FILE: KDSR-classic/template.py
function set_template (line 1) | def set_template(args):
FILE: KDSR-classic/trainer_anisonoise_stage3.py
class Trainer (line 15) | class Trainer():
method __init__ (line 16) | def __init__(self, args, loader, model_TA, my_loss, ckp):
method train (line 36) | def train(self):
method test (line 118) | def test(self):
method crop_border (line 193) | def crop_border(self, img_hr, scale):
method get_patch (line 200) | def get_patch(self, img, patch_size=48, scale=4):
method crop (line 208) | def crop(self, img_hr):
method terminate (line 216) | def terminate(self):
FILE: KDSR-classic/trainer_anisonoise_stage4.py
class Trainer (line 11) | class Trainer():
method __init__ (line 12) | def __init__(self, args, loader, model_ST, model_TA, my_loss, ckp):
method train (line 37) | def train(self):
method test (line 152) | def test(self):
method crop_border (line 222) | def crop_border(self, img_hr, scale):
method terminate (line 229) | def terminate(self):
FILE: KDSR-classic/trainer_iso_stage3.py
class Trainer (line 16) | class Trainer():
method __init__ (line 17) | def __init__(self, args, loader, model_TA, my_loss, ckp):
method train (line 37) | def train(self):
method test (line 119) | def test(self):
method crop_border (line 197) | def crop_border(self, img_hr, scale):
method terminate (line 204) | def terminate(self):
FILE: KDSR-classic/trainer_iso_stage4.py
class Trainer (line 11) | class Trainer():
method __init__ (line 12) | def __init__(self, args, loader, model_ST, model_TA, my_loss, ckp):
method train (line 38) | def train(self):
method test (line 153) | def test(self):
method crop_border (line 230) | def crop_border(self, img_hr, scale):
method terminate (line 237) | def terminate(self):
FILE: KDSR-classic/utility.py
class AverageMeter (line 15) | class AverageMeter(object):
method __init__ (line 17) | def __init__(self):
method reset (line 20) | def reset(self):
method update (line 26) | def update(self, val, n=1):
class timer (line 33) | class timer():
method __init__ (line 34) | def __init__(self):
method tic (line 38) | def tic(self):
method toc (line 41) | def toc(self):
method hold (line 44) | def hold(self):
method release (line 47) | def release(self):
method reset (line 53) | def reset(self):
class checkpoint (line 57) | class checkpoint():
method __init__ (line 58) | def __init__(self, args):
method save (line 85) | def save(self, trainer, epoch, is_best=False):
method add_log (line 97) | def add_log(self, log):
method write_log (line 100) | def write_log(self, log, refresh=False):
method done (line 107) | def done(self):
method plot_psnr (line 110) | def plot_psnr(self, epoch):
method save_results (line 128) | def save_results(self, filename, save_list, scale):
function quantize (line 136) | def quantize(img, rgb_range):
function calc_psnr (line 141) | def calc_psnr(sr, hr, scale, rgb_range, benchmark=False):
function calc_ssim (line 162) | def calc_ssim(img1, img2, scale=2, benchmark=False):
function ssim (line 199) | def ssim(img1, img2):
function make_optimizer (line 222) | def make_optimizer(args, my_model):
function make_scheduler (line 243) | def make_scheduler(args, my_optimizer):
FILE: KDSR-classic/utility2.py
class AverageMeter (line 14) | class AverageMeter(object):
method __init__ (line 16) | def __init__(self):
method reset (line 19) | def reset(self):
method update (line 25) | def update(self, val, n=1):
class timer (line 32) | class timer():
method __init__ (line 33) | def __init__(self):
method tic (line 37) | def tic(self):
method toc (line 40) | def toc(self):
method hold (line 43) | def hold(self):
method release (line 46) | def release(self):
method reset (line 52) | def reset(self):
class checkpoint (line 56) | class checkpoint():
method __init__ (line 57) | def __init__(self, args):
method save (line 83) | def save(self, trainer, epoch, is_best=False):
method add_log (line 95) | def add_log(self, log):
method write_log (line 98) | def write_log(self, log, refresh=False):
method done (line 105) | def done(self):
method plot_psnr (line 108) | def plot_psnr(self, epoch):
method save_results (line 126) | def save_results(self, filename, save_list, scale):
function quantize (line 134) | def quantize(img, rgb_range):
function calc_psnr (line 139) | def calc_psnr(sr, hr, scale, rgb_range, benchmark=False):
function calc_ssim (line 160) | def calc_ssim(img1, img2, scale=2, benchmark=False):
function ssim (line 197) | def ssim(img1, img2):
function make_optimizer (line 220) | def make_optimizer(args, my_model):
function make_scheduler (line 241) | def make_scheduler(args, my_optimizer):
FILE: KDSR-classic/utility3.py
function tensor2img (line 13) | def tensor2img(tensor, out_type=np.float32, min_max=(0, 1)):
function bgr2ycbcr (line 42) | def bgr2ycbcr(img, only_y=True):
function calc_psnr (line 66) | def calc_psnr(img1, img2,scale):
function ssim (line 82) | def ssim(img1, img2):
function calculate_ssim (line 106) | def calculate_ssim(img1, img2):
FILE: KDSR-classic/utils/util.py
function cal_sigma (line 9) | def cal_sigma(sig_x, sig_y, radians):
function anisotropic_gaussian_kernel (line 22) | def anisotropic_gaussian_kernel(batch, kernel_size, covar):
function isotropic_gaussian_kernel (line 35) | def isotropic_gaussian_kernel(batch, kernel_size, sigma):
function random_anisotropic_gaussian_kernel (line 44) | def random_anisotropic_gaussian_kernel(batch=1, kernel_size=21, lambda_m...
function stable_anisotropic_gaussian_kernel (line 54) | def stable_anisotropic_gaussian_kernel(kernel_size=21, theta=0, lambda_1...
function random_isotropic_gaussian_kernel (line 64) | def random_isotropic_gaussian_kernel(batch=1, kernel_size=21, sig_min=0....
function stable_isotropic_gaussian_kernel (line 70) | def stable_isotropic_gaussian_kernel(kernel_size=21, sig=4.0):
function random_gaussian_kernel (line 76) | def random_gaussian_kernel(batch, kernel_size=21, blur_type='iso_gaussia...
function stable_gaussian_kernel (line 83) | def stable_gaussian_kernel(kernel_size=21, blur_type='iso_gaussian', sig...
class bicubic (line 91) | class bicubic(nn.Module):
method __init__ (line 92) | def __init__(self):
method cubic (line 95) | def cubic(self, x):
method contribute (line 106) | def contribute(self, in_size, out_size, scale):
method forward (line 151) | def forward(self, input, scale=1/4):
class Gaussin_Kernel (line 171) | class Gaussin_Kernel(object):
method __init__ (line 172) | def __init__(self, kernel_size=21, blur_type='iso_gaussian',
method __call__ (line 188) | def __call__(self, batch, random):
class BatchBlur (line 201) | class BatchBlur(nn.Module):
method __init__ (line 202) | def __init__(self, kernel_size=21):
method forward (line 210) | def forward(self, input, kernel):
class SRMDPreprocessing (line 228) | class SRMDPreprocessing(object):
method __init__ (line 229) | def __init__(self,
method __call__ (line 274) | def __call__(self, hr_tensor, random=True):
class BicubicPreprocessing (line 313) | class BicubicPreprocessing(object):
method __init__ (line 314) | def __init__(self,
method __call__ (line 337) | def __call__(self, hr_tensor, random=True):
FILE: KDSR-classic/utils/util2.py
function cal_sigma (line 9) | def cal_sigma(sig_x, sig_y, radians):
function anisotropic_gaussian_kernel (line 22) | def anisotropic_gaussian_kernel(batch, kernel_size, covar):
function isotropic_gaussian_kernel (line 35) | def isotropic_gaussian_kernel(batch, kernel_size, sigma):
function random_anisotropic_gaussian_kernel (line 44) | def random_anisotropic_gaussian_kernel(batch=1, kernel_size=21, lambda_m...
function stable_anisotropic_gaussian_kernel (line 54) | def stable_anisotropic_gaussian_kernel(kernel_size=21, theta=0, lambda_1...
function random_isotropic_gaussian_kernel (line 64) | def random_isotropic_gaussian_kernel(batch=1, kernel_size=21, sig_min=0....
function stable_isotropic_gaussian_kernel (line 70) | def stable_isotropic_gaussian_kernel(kernel_size=21, sig=4.0):
function random_gaussian_kernel (line 76) | def random_gaussian_kernel(batch, kernel_size=21, blur_type='iso_gaussia...
function stable_gaussian_kernel (line 83) | def stable_gaussian_kernel(kernel_size=21, blur_type='iso_gaussian', sig...
class bicubic (line 91) | class bicubic(nn.Module):
method __init__ (line 92) | def __init__(self):
method cubic (line 95) | def cubic(self, x):
method contribute (line 106) | def contribute(self, in_size, out_size, scale):
method forward (line 151) | def forward(self, input, scale=1/4):
class Gaussin_Kernel (line 171) | class Gaussin_Kernel(object):
method __init__ (line 172) | def __init__(self, kernel_size=21, blur_type='iso_gaussian',
method __call__ (line 188) | def __call__(self, batch, random):
class BatchBlur (line 201) | class BatchBlur(nn.Module):
method __init__ (line 202) | def __init__(self, kernel_size=21):
method forward (line 210) | def forward(self, input, kernel):
class SRMDPreprocessing (line 228) | class SRMDPreprocessing(object):
method __init__ (line 229) | def __init__(self,
method __call__ (line 274) | def __call__(self, hr_tensor, random=True):
Condensed preview — 135 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (460K chars).
[
{
"path": "KDSR-GAN/.gitignore",
"chars": 1912,
"preview": "# ignored folders\ndatasets/*\nexperiments/*\nresults/*\ntb_logger/*\nwandb/*\ntmp/*\nrealesrgan/weights/*\n\nversion.py\n\n# Byte-"
},
{
"path": "KDSR-GAN/.pre-commit-config.yaml",
"chars": 1483,
"preview": "repos:\n # flake8\n - repo: https://github.com/PyCQA/flake8\n rev: 3.8.3\n hooks:\n - id: flake8\n args: ["
},
{
"path": "KDSR-GAN/CODE_OF_CONDUCT.md",
"chars": 5251,
"preview": "# Contributor Covenant Code of Conduct\n\n## Our Pledge\n\nWe as members, contributors, and leaders pledge to make participa"
},
{
"path": "KDSR-GAN/Metric/DISTS/DISTS_pytorch/DISTS_pt.py",
"chars": 8604,
"preview": "# This is a pytoch implementation of DISTS metric.\n# Requirements: python >= 3.6, pytorch >= 1.0\n\nimport numpy as np\nimp"
},
{
"path": "KDSR-GAN/Metric/DISTS/DISTS_tensorflow/DISTS_tf.py",
"chars": 5335,
"preview": "# This is a tensorflow implementation of DISTS metric.\n# Requirements: python >= 3.6, tensorflow-gpu >= 1.15\n\nimport ten"
},
{
"path": "KDSR-GAN/Metric/DISTS/LICENSE",
"chars": 1067,
"preview": "MIT License\n\nCopyright (c) 2020 Keyan Ding\n\nPermission is hereby granted, free of charge, to any person obtaining a copy"
},
{
"path": "KDSR-GAN/Metric/DISTS/requirements.txt",
"chars": 10,
"preview": "torch>=1.0"
},
{
"path": "KDSR-GAN/Metric/LPIPS.py",
"chars": 1639,
"preview": "import cv2\nimport glob\nimport numpy as np\nimport os.path as osp\nfrom torchvision.transforms.functional import normalize\n"
},
{
"path": "KDSR-GAN/Metric/PSNR.py",
"chars": 1399,
"preview": "import cv2\nimport glob\nimport numpy as np\nimport os.path as osp\nfrom torchvision.transforms.functional import normalize\n"
},
{
"path": "KDSR-GAN/Metric/dists.py",
"chars": 6169,
"preview": "# This is a pytoch implementation of DISTS metric.\n# Requirements: python >= 3.6, pytorch >= 1.0\n\nimport numpy as np\nimp"
},
{
"path": "KDSR-GAN/Metric/pip.sh",
"chars": 184,
"preview": "pip install lpips --index-url http://pypi.douban.com/simple --trusted-host pypi.douban.com\npip install basicsr --index-u"
},
{
"path": "KDSR-GAN/README.md",
"chars": 3596,
"preview": "# KDSR-GAN\n\nThis project is the official implementation of 'Knowledge Distillation based Degradation Estimation for Blin"
},
{
"path": "KDSR-GAN/VERSION",
"chars": 8,
"preview": "0.2.5.0\n"
},
{
"path": "KDSR-GAN/docs/CONTRIBUTING.md",
"chars": 1876,
"preview": "# Contributing to Real-ESRGAN\n\n:art: Real-ESRGAN needs your contributions. Any contributions are welcome, such as new fe"
},
{
"path": "KDSR-GAN/docs/FAQ.md",
"chars": 736,
"preview": "# FAQ\n\n1. **Q: How to select models?**<br>\nA: Please refer to [docs/model_zoo.md](docs/model_zoo.md)\n\n1. **Q: Can `face_"
},
{
"path": "KDSR-GAN/docs/Training.md",
"chars": 11313,
"preview": "# :computer: How to Train/Finetune Real-ESRGAN\n\n- [Train Real-ESRGAN](#train-real-esrgan)\n - [Overview](#overview)\n - "
},
{
"path": "KDSR-GAN/docs/Training_CN.md",
"chars": 8532,
"preview": "# :computer: 如何训练/微调 Real-ESRGAN\n\n- [训练 Real-ESRGAN](#训练-real-esrgan)\n - [概述](#概述)\n - [准备数据集](#准备数据集)\n - [训练 Real-ESR"
},
{
"path": "KDSR-GAN/docs/anime_comparisons.md",
"chars": 8181,
"preview": "# Comparisons among different anime models\n\n[English](anime_comparisons.md) **|** [简体中文](anime_comparisons_CN.md)\n\n## Up"
},
{
"path": "KDSR-GAN/docs/anime_comparisons_CN.md",
"chars": 7560,
"preview": "# 动漫视频模型比较\n\n[English](anime_comparisons.md) **|** [简体中文](anime_comparisons_CN.md)\n\n## 更新\n\n- 2022/04/24: 发布 **AnimeVideo-"
},
{
"path": "KDSR-GAN/docs/anime_model.md",
"chars": 3274,
"preview": "# Anime Model\n\n:white_check_mark: We add [*RealESRGAN_x4plus_anime_6B.pth*](https://github.com/xinntao/Real-ESRGAN/relea"
},
{
"path": "KDSR-GAN/docs/anime_video_model.md",
"chars": 5727,
"preview": "# Anime Video Models\n\n:white_check_mark: We add small models that are optimized for anime videos :-)<br>\nMore comparison"
},
{
"path": "KDSR-GAN/docs/feedback.md",
"chars": 590,
"preview": "# Feedback 反馈\n\n## 动漫插画模型\n\n1. 视频处理不了: 目前的模型，不是针对视频的，所以视频效果很很不好。我们在探究针对视频的模型了\n1. 景深虚化有问题: 现在的模型把一些景深 和 特意的虚化 都复原了，感觉不好。这个后"
},
{
"path": "KDSR-GAN/docs/model_zoo.md",
"chars": 3922,
"preview": "# :european_castle: Model Zoo\n\n- [For General Images](#for-general-images)\n- [For Anime Images](#for-anime-images)\n- [Fo"
},
{
"path": "KDSR-GAN/docs/ncnn_conversion.md",
"chars": 580,
"preview": "# Instructions on converting to NCNN models\n\n1. Convert to onnx model with `scripts/pytorch2onnx.py`. Remember to modify"
},
{
"path": "KDSR-GAN/kdsrgan/__init__.py",
"chars": 144,
"preview": "# flake8: noqa\nfrom .losses import *\nfrom .archs import *\nfrom .data import *\nfrom .models import *\nfrom .utils import *"
},
{
"path": "KDSR-GAN/kdsrgan/archs/ST_arch.py",
"chars": 6030,
"preview": "from basicsr.utils.registry import ARCH_REGISTRY\nfrom torch.nn import functional as F\nimport torch\nfrom torch import nn\n"
},
{
"path": "KDSR-GAN/kdsrgan/archs/TA_arch.py",
"chars": 6285,
"preview": "from basicsr.utils.registry import ARCH_REGISTRY\nfrom torch.nn import functional as F\nimport torch\nfrom torch import nn\n"
},
{
"path": "KDSR-GAN/kdsrgan/archs/__init__.py",
"chars": 497,
"preview": "import importlib\nfrom basicsr.utils import scandir\nfrom os import path as osp\n\n# automatically scan and import arch modu"
},
{
"path": "KDSR-GAN/kdsrgan/archs/common.py",
"chars": 1976,
"preview": "import math\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef default_conv(in_channels, out_chann"
},
{
"path": "KDSR-GAN/kdsrgan/archs/discriminator_arch.py",
"chars": 3020,
"preview": "from basicsr.utils.registry import ARCH_REGISTRY\nfrom torch import nn as nn\nfrom torch.nn import functional as F\nfrom to"
},
{
"path": "KDSR-GAN/kdsrgan/archs/srvgg_arch.py",
"chars": 2722,
"preview": "from basicsr.utils.registry import ARCH_REGISTRY\nfrom torch import nn as nn\nfrom torch.nn import functional as F\n\n\n@ARCH"
},
{
"path": "KDSR-GAN/kdsrgan/data/__init__.py",
"chars": 516,
"preview": "import importlib\nfrom basicsr.utils import scandir\nfrom os import path as osp\n\n# automatically scan and import dataset m"
},
{
"path": "KDSR-GAN/kdsrgan/data/kdsrgan_dataset.py",
"chars": 8719,
"preview": "import cv2\nimport math\nimport numpy as np\nimport os\nimport os.path as osp\nimport random\nimport time\nimport torch\nfrom ba"
},
{
"path": "KDSR-GAN/kdsrgan/data/kdsrgan_paired_dataset.py",
"chars": 5009,
"preview": "import os\nfrom basicsr.data.data_util import paired_paths_from_folder, paired_paths_from_lmdb\nfrom basicsr.data.transfor"
},
{
"path": "KDSR-GAN/kdsrgan/losses/__init__.py",
"chars": 498,
"preview": "import importlib\nfrom basicsr.utils import scandir\nfrom os import path as osp\n\n# automatically scan and import arch modu"
},
{
"path": "KDSR-GAN/kdsrgan/losses/my_loss.py",
"chars": 1297,
"preview": "import torch\nfrom torch import nn as nn\nfrom torch.nn import functional as F\nfrom basicsr.utils.registry import LOSS_REG"
},
{
"path": "KDSR-GAN/kdsrgan/models/__init__.py",
"chars": 507,
"preview": "import importlib\nfrom basicsr.utils import scandir\nfrom os import path as osp\n\n# automatically scan and import model mod"
},
{
"path": "KDSR-GAN/kdsrgan/models/kdsrgan_ST_model.py",
"chars": 14035,
"preview": "import numpy as np\nimport random\nimport torch\nfrom basicsr.data.degradations import random_add_gaussian_noise_pt, random"
},
{
"path": "KDSR-GAN/kdsrgan/models/kdsrgan_TA_model.py",
"chars": 11339,
"preview": "import numpy as np\nimport random\nimport torch\nfrom basicsr.data.degradations import random_add_gaussian_noise_pt, random"
},
{
"path": "KDSR-GAN/kdsrgan/models/kdsrnet_ST_model.py",
"chars": 14556,
"preview": "import numpy as np\nimport random\nimport torch\nfrom basicsr.data.degradations import random_add_gaussian_noise_pt, random"
},
{
"path": "KDSR-GAN/kdsrgan/models/kdsrnet_TA_model.py",
"chars": 10442,
"preview": "import numpy as np\nimport random\nimport torch\nfrom basicsr.data.degradations import random_add_gaussian_noise_pt, random"
},
{
"path": "KDSR-GAN/kdsrgan/test.py",
"chars": 269,
"preview": "# flake8: noqa\nimport os.path as osp\nfrom basicsr.test import test_pipeline\n\nimport kdsrgan.archs\nimport kdsrgan.data\nim"
},
{
"path": "KDSR-GAN/kdsrgan/train.py",
"chars": 345,
"preview": "# flake8: noqa\nimport os.path as osp\nfrom basicsr.train import train_pipeline\n\nimport kdsrgan.archs\nimport kdsrgan.data\n"
},
{
"path": "KDSR-GAN/kdsrgan/utils.py",
"chars": 11827,
"preview": "import cv2\nimport math\nimport numpy as np\nimport os\nimport queue\nimport threading\nimport torch\nfrom basicsr.utils.downlo"
},
{
"path": "KDSR-GAN/kdsrgan/weights/README.md",
"chars": 54,
"preview": "# Weights\n\nPut the downloaded weights to this folder.\n"
},
{
"path": "KDSR-GAN/options/test_kdsrgan_x4ST.yml",
"chars": 1429,
"preview": "# general settings\nname: test_KDSRGANx4STplus_400k_B12G4\nmodel_type: KDSRGANSTModel\nscale: 4\nnum_gpu: auto # auto: can "
},
{
"path": "KDSR-GAN/options/test_kdsrnet_x4TA.yml",
"chars": 1221,
"preview": "# general settings\nname: train_KDSRNetTAx4plus_1000k_B12G4\nmodel_type: KDSRNetTAModel\nscale: 4\nnum_gpu: auto # auto: ca"
},
{
"path": "KDSR-GAN/options/train_kdsrgan_x4ST.yml",
"chars": 4819,
"preview": "# general settings\nname: train_KDSRGANx4STplus_400k_B12G4\nmodel_type: KDSRGANSTModel\nscale: 4\nnum_gpu: auto # auto: can"
},
{
"path": "KDSR-GAN/options/train_kdsrgan_x4STV2.yml",
"chars": 4814,
"preview": "# general settings\nname: train_KDSRGANx4STplus_400k_V2\nmodel_type: KDSRGANSTModel\nscale: 4\nnum_gpu: auto # auto: can in"
},
{
"path": "KDSR-GAN/options/train_kdsrnet_x4ST.yml",
"chars": 3588,
"preview": "# general settings\nname: train_KDSRNetSTx4_1000k\nmodel_type: KDSRNetSTModel\nscale: 4\nnum_gpu: auto # auto: can infer fr"
},
{
"path": "KDSR-GAN/options/train_kdsrnet_x4TA.yml",
"chars": 3142,
"preview": "# general settings\nname: train_KDSRNetTAx4plus_1000k_B12G4\nmodel_type: KDSRNetTAModel\nscale: 4\nnum_gpu: auto # auto: ca"
},
{
"path": "KDSR-GAN/pip.sh",
"chars": 103,
"preview": "pip install basicsr \npip install -r requirements.txt\npip install pandas \nsudo python3 setup.py develop\n"
},
{
"path": "KDSR-GAN/requirements.txt",
"chars": 105,
"preview": "basicsr>=1.3.3.11\nfacexlib>=0.2.0.3\ngfpgan>=0.2.1\nnumpy\nopencv-python\nPillow\ntorch>=1.7\ntorchvision\ntqdm\n"
},
{
"path": "KDSR-GAN/scripts/extract_subimages.py",
"chars": 5116,
"preview": "import argparse\nimport cv2\nimport numpy as np\nimport os\nimport sys\nfrom basicsr.utils import scandir\nfrom multiprocessin"
},
{
"path": "KDSR-GAN/scripts/extract_subimages_DF2K.py",
"chars": 5092,
"preview": "import argparse\nimport cv2\nimport numpy as np\nimport os\nimport sys\nfrom basicsr.utils import scandir\nfrom multiprocessin"
},
{
"path": "KDSR-GAN/scripts/generate_meta_info.py",
"chars": 2114,
"preview": "import argparse\nimport cv2\nimport glob\nimport os\n\n\ndef main(args):\n txt_file = open(args.meta_info, 'w')\n for fold"
},
{
"path": "KDSR-GAN/scripts/generate_meta_info_DF2K.py",
"chars": 2221,
"preview": "import argparse\nimport cv2\nimport glob\nimport os\n\n\ndef main(args):\n txt_file = open(args.meta_info, 'w')\n for fold"
},
{
"path": "KDSR-GAN/scripts/generate_meta_info_OST.py",
"chars": 2118,
"preview": "import argparse\nimport cv2\nimport glob\nimport os\n\n\ndef main(args):\n txt_file = open(args.meta_info, 'w')\n for fold"
},
{
"path": "KDSR-GAN/scripts/generate_meta_info_pairdata.py",
"chars": 1987,
"preview": "import argparse\nimport glob\nimport os\n\n\ndef main(args):\n txt_file = open(args.meta_info, 'w')\n # sca images\n im"
},
{
"path": "KDSR-GAN/scripts/generate_multiscale_DF2K.py",
"chars": 1777,
"preview": "import argparse\nimport glob\nimport os\nfrom PIL import Image\n\n\ndef main(args):\n # For DF2K, we consider the following "
},
{
"path": "KDSR-GAN/scripts/pytorch2onnx.py",
"chars": 1247,
"preview": "import argparse\nimport torch\nimport torch.onnx\nfrom basicsr.archs.rrdbnet_arch import RRDBNet\n\n\ndef main(args):\n # An"
},
{
"path": "KDSR-GAN/setup.cfg",
"chars": 684,
"preview": "[flake8]\nignore =\n # line break before binary operator (W503)\n W503,\n # line break after binary operator (W504)"
},
{
"path": "KDSR-GAN/setup.py",
"chars": 3142,
"preview": "#!/usr/bin/env python\n\nfrom setuptools import find_packages, setup\n\nimport os\nimport subprocess\nimport time\n\nversion_fil"
},
{
"path": "KDSR-GAN/test.sh",
"chars": 65,
"preview": "\n\n\n\npython3 kdsrgan/test.py -opt options/test_kdsrgan_x4ST.yml "
},
{
"path": "KDSR-GAN/tests/data/gt.lmdb/meta_info.txt",
"chars": 49,
"preview": "baboon.png (480,500,3) 1\ncomic.png (360,240,3) 1\n"
},
{
"path": "KDSR-GAN/tests/data/lq.lmdb/meta_info.txt",
"chars": 47,
"preview": "baboon.png (120,125,3) 1\ncomic.png (80,60,3) 1\n"
},
{
"path": "KDSR-GAN/tests/data/meta_info_gt.txt",
"chars": 21,
"preview": "baboon.png\ncomic.png\n"
},
{
"path": "KDSR-GAN/tests/data/meta_info_pair.txt",
"chars": 56,
"preview": "gt/baboon.png, lq/baboon.png\ngt/comic.png, lq/comic.png\n"
},
{
"path": "KDSR-GAN/tests/data/test_realesrgan_dataset.yml",
"chars": 700,
"preview": "name: Demo\ntype: RealESRGANDataset\ndataroot_gt: tests/data/gt\nmeta_info: tests/data/meta_info_gt.txt\nio_backend:\n type:"
},
{
"path": "KDSR-GAN/tests/data/test_realesrgan_model.yml",
"chars": 2101,
"preview": "scale: 4\nnum_gpu: 1\nmanual_seed: 0\nis_train: True\ndist: False\n\n# ----------------- options for synthesizing training dat"
},
{
"path": "KDSR-GAN/tests/data/test_realesrgan_paired_dataset.yml",
"chars": 222,
"preview": "name: Demo\ntype: RealESRGANPairedDataset\nscale: 4\ndataroot_gt: tests/data\ndataroot_lq: tests/data\nmeta_info: tests/data/"
},
{
"path": "KDSR-GAN/tests/data/test_realesrnet_model.yml",
"chars": 1330,
"preview": "scale: 4\nnum_gpu: 1\nmanual_seed: 0\nis_train: True\ndist: False\n\n# ----------------- options for synthesizing training dat"
},
{
"path": "KDSR-GAN/tests/test_dataset.py",
"chars": 6037,
"preview": "import pytest\nimport yaml\n\nfrom realesrgan.data.realesrgan_dataset import RealESRGANDataset\nfrom realesrgan.data.realesr"
},
{
"path": "KDSR-GAN/tests/test_discriminator_arch.py",
"chars": 561,
"preview": "import torch\n\nfrom realesrgan.archs.discriminator_arch import UNetDiscriminatorSN\n\n\ndef test_unetdiscriminatorsn():\n "
},
{
"path": "KDSR-GAN/tests/test_model.py",
"chars": 4861,
"preview": "import torch\nimport yaml\nfrom basicsr.archs.rrdbnet_arch import RRDBNet\nfrom basicsr.data.paired_image_dataset import Pa"
},
{
"path": "KDSR-GAN/tests/test_utils.py",
"chars": 3089,
"preview": "import numpy as np\nfrom basicsr.archs.rrdbnet_arch import RRDBNet\n\nfrom realesrgan.utils import RealESRGANer\n\n\ndef test_"
},
{
"path": "KDSR-GAN/train_KDSRGAN_S.sh",
"chars": 188,
"preview": "\nCUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \\\npython3 -m torch.distributed.launch --nproc_per_node=8 --master_port=4397 kdsrga"
},
{
"path": "KDSR-GAN/train_KDSRNet_S.sh",
"chars": 190,
"preview": "\n\n\n\nCUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \\\npython3 -m torch.distributed.launch --nproc_per_node=8 --master_port=4349 kd"
},
{
"path": "KDSR-GAN/train_KDSRNet_T.sh",
"chars": 189,
"preview": "\n\n\n\nCUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \\\npython3 -m torch.distributed.launch --nproc_per_node=8 --master_port=3309 kds"
},
{
"path": "KDSR-classic/README.md",
"chars": 2843,
"preview": "# KDSR-classic\n\n\nThis project is the official implementation of 'Knowledge Distillation based Degradation Estimation for"
},
{
"path": "KDSR-classic/cal_parms.py",
"chars": 467,
"preview": "def count_param(model):\n param_count = 0\n for param in model.parameters():\n param_count += param.view(-1).s"
},
{
"path": "KDSR-classic/data/__init__.py",
"chars": 1986,
"preview": "from importlib import import_module\n#from dataloader import MSDataLoader\nfrom torch.utils.data import dataloader\nfrom to"
},
{
"path": "KDSR-classic/data/benchmark.py",
"chars": 583,
"preview": "import os\nfrom data import common\nfrom data import multiscalesrdata as srdata\n\n\nclass Benchmark(srdata.SRData):\n def "
},
{
"path": "KDSR-classic/data/common.py",
"chars": 1398,
"preview": "import random\n\nimport numpy as np\nimport skimage.color as sc\n\nimport torch\n\n\ndef get_patch(hr, patch_size=48, scale=2, m"
},
{
"path": "KDSR-classic/data/common2.py",
"chars": 1780,
"preview": "import random\n\nimport numpy as np\nimport skimage.color as sc\n\nimport torch\n\n\ndef get_patch(*args, patch_size=96, scale=2"
},
{
"path": "KDSR-classic/data/df2k.py",
"chars": 970,
"preview": "import os\nfrom data import multiscalesrdata\n\n\nclass DF2K(multiscalesrdata.SRData):\n def __init__(self, args, name='DF"
},
{
"path": "KDSR-classic/data/multiscalesrdata.py",
"chars": 5147,
"preview": "import os\nimport glob\nimport random\nimport pickle\n\nfrom data import common\n\nimport numpy as np\nimport imageio\nimport tor"
},
{
"path": "KDSR-classic/data/multiscalesrdata2.py",
"chars": 5374,
"preview": "import os\nimport glob\nimport random\nimport pickle\n\nfrom data import common2 as common\n\nimport numpy as np\nimport imageio"
},
{
"path": "KDSR-classic/dataloader.py",
"chars": 4776,
"preview": "import sys\nimport threading\nimport queue\nimport random\nimport collections\n\nimport torch\nimport torch.multiprocessing as "
},
{
"path": "KDSR-classic/loss/__init__.py",
"chars": 4703,
"preview": "import os\nfrom importlib import import_module\n\nimport matplotlib.pyplot as plt\n\nimport numpy as np\n\nimport torch\nimport "
},
{
"path": "KDSR-classic/loss/adversarial.py",
"chars": 3320,
"preview": "import utility\nfrom model import common\nfrom loss import discriminator\n\nimport torch\nimport torch.nn as nn\nimport torch."
},
{
"path": "KDSR-classic/loss/discriminator.py",
"chars": 1287,
"preview": "from model import common\n\nimport torch.nn as nn\n\nclass Discriminator(nn.Module):\n def __init__(self, args, gan_type='"
},
{
"path": "KDSR-classic/loss/vgg.py",
"chars": 1093,
"preview": "from model import common\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.models a"
},
{
"path": "KDSR-classic/main_anisonoise_KDSRsMx4_stage3.sh",
"chars": 810,
"preview": "## noise-free degradations with isotropic Gaussian blurs\n# training knowledge distillation\nCUDA_VISIBLE_DEVICES=0,1,2,3 "
},
{
"path": "KDSR-classic/main_anisonoise_KDSRsMx4_stage4.sh",
"chars": 1102,
"preview": "## noise-free degradations with isotropic Gaussian blurs\n# training knowledge distillation\nCUDA_VISIBLE_DEVICES=0,1,2,3 "
},
{
"path": "KDSR-classic/main_anisonoise_stage3.py",
"chars": 780,
"preview": "from option import args\nimport torch\nimport utility\nimport data\nimport model_TA\nimport loss\nfrom trainer_anisonoise_stag"
},
{
"path": "KDSR-classic/main_anisonoise_stage4.py",
"chars": 832,
"preview": "from option import args\nimport torch\nimport utility\nimport data\nimport model_TA\nimport model_ST\nimport loss\nfrom trainer"
},
{
"path": "KDSR-classic/main_iso_KDSRsLx4_stage3.sh",
"chars": 811,
"preview": "## noise-free degradations with isotropic Gaussian blurs\n# training knowledge distillation\nCUDA_VISIBLE_DEVICES=0,1,2,3 "
},
{
"path": "KDSR-classic/main_iso_KDSRsLx4_stage4.sh",
"chars": 1113,
"preview": "## noise-free degradations with isotropic Gaussian blurs\n# training knowledge distillation\nCUDA_VISIBLE_DEVICES=0,1,2,3 "
},
{
"path": "KDSR-classic/main_iso_KDSRsMx4_stage3.sh",
"chars": 718,
"preview": "## noise-free degradations with isotropic Gaussian blurs\n# training knowledge distillation\nCUDA_VISIBLE_DEVICES=4,5,6,7 "
},
{
"path": "KDSR-classic/main_iso_KDSRsMx4_stage4.sh",
"chars": 1018,
"preview": "## noise-free degradations with isotropic Gaussian blurs\n# training knowledge distillation\nCUDA_VISIBLE_DEVICES=4,5,6,7 "
},
{
"path": "KDSR-classic/main_iso_stage3.py",
"chars": 773,
"preview": "from option import args\nimport torch\nimport utility\nimport data\nimport model_TA\nimport loss\nfrom trainer_iso_stage3 impo"
},
{
"path": "KDSR-classic/main_iso_stage4.py",
"chars": 825,
"preview": "from option import args\nimport torch\nimport utility\nimport data\nimport model_TA\nimport model_ST\nimport loss\nfrom trainer"
},
{
"path": "KDSR-classic/model/__init__.py",
"chars": 5873,
"preview": "import os\nfrom importlib import import_module\n\nimport torch\nimport torch.nn as nn\n\n\nclass Model(nn.Module):\n def __in"
},
{
"path": "KDSR-classic/model/blindsr.py",
"chars": 2078,
"preview": "import torch\nfrom torch import nn\nimport model.common as common\nimport torch.nn.functional as F\nfrom moco.builder import"
},
{
"path": "KDSR-classic/model/common.py",
"chars": 1978,
"preview": "import math\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef default_conv(in_channels, out_chann"
},
{
"path": "KDSR-classic/model_ST/__init__.py",
"chars": 5876,
"preview": "import os\nfrom importlib import import_module\n\nimport torch\nimport torch.nn as nn\n\n\nclass Model(nn.Module):\n def __in"
},
{
"path": "KDSR-classic/model_ST/blindsr.py",
"chars": 6070,
"preview": "import torch\nfrom torch import nn\nimport model.common as common\nimport torch.nn.functional as F\n\n\ndef make_model(args):\n"
},
{
"path": "KDSR-classic/model_ST/common.py",
"chars": 1359,
"preview": "import math\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef default_conv(in_channels, out_chann"
},
{
"path": "KDSR-classic/model_TA/__init__.py",
"chars": 6479,
"preview": "import os\nfrom importlib import import_module\n\nimport torch\nimport torch.nn as nn\n\n\nclass Model(nn.Module):\n def __in"
},
{
"path": "KDSR-classic/model_TA/blindsr.py",
"chars": 5999,
"preview": "import torch\nfrom torch import nn\nimport model.common as common\nimport torch.nn.functional as F\n\n\ndef make_model(args):\n"
},
{
"path": "KDSR-classic/model_TA/common.py",
"chars": 1359,
"preview": "import math\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef default_conv(in_channels, out_chann"
},
{
"path": "KDSR-classic/option.py",
"chars": 10287,
"preview": "import argparse\nimport template\n\nparser = argparse.ArgumentParser(description='EDSR and MDSR')\n\nparser.add_argument('--d"
},
{
"path": "KDSR-classic/template.py",
"chars": 1138,
"preview": "def set_template(args):\n # Set the templates here\n if args.template.find('jpeg') >= 0:\n args.data_train = '"
},
{
"path": "KDSR-classic/test.py",
"chars": 538,
"preview": "from option import args\nimport torch\nimport utility\nimport data\nimport model\nimport loss\nfrom trainer import Trainer\n\n\ni"
},
{
"path": "KDSR-classic/test_anisonoise_KDSRsMx4.sh",
"chars": 565,
"preview": "CUDA_VISIBLE_DEVICES=4 python3 main_anisonoise_stage4.py --dir_data='/root/datasets' \\\n --model='blindsr' "
},
{
"path": "KDSR-classic/test_iso_KDSRsLx4.sh",
"chars": 612,
"preview": "CUDA_VISIBLE_DEVICES=4 python3 main_iso_stage4.py --test_only \\\n --dir_data='/root/datasets' \\\n "
},
{
"path": "KDSR-classic/test_iso_KDSRsMx4.sh",
"chars": 520,
"preview": "CUDA_VISIBLE_DEVICES=4 python3 main_iso_stage4.py --test_only \\\n --dir_data='/root/datasets' \\\n "
},
{
"path": "KDSR-classic/trainer_anisonoise_stage3.py",
"chars": 8048,
"preview": "import os\nimport utility2\nimport torch\nfrom decimal import Decimal\nimport torch.nn.functional as F\nfrom utils import uti"
},
{
"path": "KDSR-classic/trainer_anisonoise_stage4.py",
"chars": 10100,
"preview": "import os\nimport utility\nimport torch\nfrom decimal import Decimal\nimport torch.nn.functional as F\nfrom utils import util"
},
{
"path": "KDSR-classic/trainer_iso_stage3.py",
"chars": 8124,
"preview": "import os\nimport utility2\nimport utility\nimport torch\nfrom decimal import Decimal\nimport torch.nn.functional as F\nfrom u"
},
{
"path": "KDSR-classic/trainer_iso_stage4.py",
"chars": 10654,
"preview": "import os\nimport utility\nimport torch\nfrom decimal import Decimal\nimport torch.nn.functional as F\nfrom utils import util"
},
{
"path": "KDSR-classic/utility.py",
"chars": 7912,
"preview": "import os\nimport math\nimport time\nimport datetime\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport scipy.misc a"
},
{
"path": "KDSR-classic/utility2.py",
"chars": 7852,
"preview": "import os\nimport math\nimport time\nimport datetime\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport scipy.misc a"
},
{
"path": "KDSR-classic/utility3.py",
"chars": 4086,
"preview": "import os\nimport math\nimport time\nimport datetime\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport scipy.misc a"
},
{
"path": "KDSR-classic/utils/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "KDSR-classic/utils/util.py",
"chars": 14078,
"preview": "import math\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport utility\n\n\ndef c"
},
{
"path": "KDSR-classic/utils/util2.py",
"chars": 12627,
"preview": "import math\nimport numpy as np\nimport utility\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef c"
},
{
"path": "README.md",
"chars": 3689,
"preview": "# Knowledge Distillation based Degradation Estimation for Blind Super-Resolution （ICLR2023）\n\n\n[Paper](https://arxiv.org/"
}
]
// ... and 6 more files (download for full content)
About this extraction
This page contains the full source code of the Zj-BinXia/KDSR GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 135 files (425.7 KB), approximately 121.2k tokens, and a symbol index with 447 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.