没研究过也更没做过超分, 但是项目又得用超分, 有人会选择BasicSR,
这边还是选择的MMEditing, 商汤的框架用惯了, 懒得写的时候直接就用。
网上的介绍比较少, 官方1.x的文档更新的也很有限。稍微写一下, 我这边只做了图像超分所以也只能写图像超分。
请注意, 该帮助仅支持0.x版本, 且似乎不支持MMagic。
图像超分顾名思义, 你希望将一张分辨率较低的图像送进网络去获得一张高分辨率的图片。
因此在做数据的时候, 需要高分辨率的原图和低分辨率的对应图片。MMEditing在超分这块仅支持DIV2K数据集, 且为修改过的DIV2K数据集。
整个数据集分为三块: 本身和下采样2 3 4倍的对应图像, Set5和Set14的两个验证集。请注意,
Set5 和 Set14数据集也是超分的Benchmark之一, 下载地址请见这里。不过最终发现不要这个也行。
整个数据集文件夹最终的呈现效果如下方所示。请注意: 所有文件必须为png格式, 且不同文件夹下对应图片名称相同。
data
├─DIV2K
│ ├─DIV2K_train_HR <--- 原分辨率的图片, 训练集
│ ├─DIV2K_train_LR_bicubic
│ │ ├─X2 <--- 2倍下采样
│ │ ├─X3 <--- 3倍下采样
│ │ └─X4 <--- 4倍下采样
│ ├─DIV2K_valid_HR <--- 原分辨率的图片, 验证集
│ └─DIV2K_valid_LR_bicubic
│ ├─X2
│ ├─X3
│ └─X4
├─val_set14 <--- 更改名字即可
│ ├─Set14_mod12
│ ├─original
│ ├─Set14_bicLRx2
│ ├─Set14_bicLRx3
│ └─Set14_bicLRx4
└─val_set5 <--- 更改名字即可
├─Set5_mod12
├─original
├─Set5_bicLRx2
├─Set5_bicLRx3
└─Set5_bicLRx4
第一步当然是要做这个数据集。
下采样的方法其实从文件夹命名就可以看出来, 一个bicubic就能搞定。
代码参照这里并做了下修改
import os
import argparse
import cv2
parser = argparse.ArgumentParser(description='Downsize images at 2x using bicubic interpolation')
parser.add_argument("-k", "--keepdims", help="keep original image dimensions in downsampled images", action="store_true")
parser.add_argument('--hr_img_dir', type=str, default=None,
help='path to high resolution image dir')
parser.add_argument('--lr_img_dir', type=str, default=None,
help='path to desired output dir for downsampled images')
args = parser.parse_args()
hr_image_dir = args.hr_img_dir
lr_image_dir = args.lr_img_dir
os.makedirs(lr_image_dir + "/X2", exist_ok=True)
os.makedirs(lr_image_dir + "/X3", exist_ok=True)
os.makedirs(lr_image_dir + "/X4", exist_ok=True)
supported_img_formats = (".bmp", ".dib", ".jpeg", ".jpg", ".jpe", ".jp2",
".png", ".pbm", ".pgm", ".ppm", ".sr", ".ras", ".tif",
".tiff")
for filename in os.listdir(hr_image_dir):
if not filename.endswith(supported_img_formats):
continue
name, ext = os.path.splitext(filename)
hr_img = cv2.imread(os.path.join(hr_image_dir, filename))
hr_img_dims = (hr_img.shape[1], hr_img.shape[0])
hr_img = cv2.GaussianBlur(hr_img, (0,0), 1, 1)
lr_image_2x = cv2.resize(hr_img, (0,0), fx=0.5, fy=0.5, interpolation=cv2.INTER_CUBIC)
if args.keepdims:
lr_image_2x = cv2.resize(lr_image_2x, hr_img_dims, interpolation=cv2.INTER_CUBIC)
cv2.imwrite(os.path.join(lr_image_dir + "/X2", filename.split('.')[0] + ext), lr_image_2x)
lr_img_3x = cv2.resize(hr_img, (0, 0), fx=(1 / 3), fy=(1 / 3),
interpolation=cv2.INTER_CUBIC)
if args.keepdims:
lr_img_3x = cv2.resize(lr_img_3x, hr_img_dims,
interpolation=cv2.INTER_CUBIC)
cv2.imwrite(os.path.join(lr_image_dir + "/X3", filename.split('.')[0] + ext), lr_img_3x)
lr_img_4x = cv2.resize(hr_img, (0, 0), fx=0.25, fy=0.25,
interpolation=cv2.INTER_CUBIC)
if args.keepdims:
lr_img_4x = cv2.resize(lr_img_4x, hr_img_dims,
interpolation=cv2.INTER_CUBIC)
cv2.imwrite(os.path.join(lr_image_dir + "/X4", filename.split('.')[0] + ext), lr_img_4x)
第二步在MMEditing内预处理数据集
官方提供的数据集格式有 SRAnnotationDataset --> 以图片文件和独立的ann.txt文件组成;
SRFolderDataset --> 直接以图片文件组成;
SRLmdbDataset --> 以图片和lmdb文件组成。官方建议的是使用lmdb数据集, io处理速度快, 但是以文件夹为核心的数据集也不错, 起码制作起来简单。
为了扩充数据 (和获取ann.txt及lmdb文件), 首先得进一遍 tools/data/super-resolution/div2k/preprocess_div2k_dataset.py。
具体命令为 (假设你的数据文件夹名为data, 如同上方所示):
python tools/data/super-resolution/div2k/preprocess_div2k_dataset.py --data-root ./data/DIV2K --make-lmdb
之后你的data文件夹多了几个lmdb和_sub文件夹, 如下方所示。_sub文件夹内即为经裁剪后的图片, lmdb就是所需的数据文件了。
data
├── DIV2K
│ ├── DIV2K_train_HR
│ ├── DIV2K_train_HR_sub
│ ├── DIV2K_train_HR_sub.lmdb
│ │ ├── data.mdb
│ │ ├── lock.mdb
│ │ ├── meta_info.txt
│ ├── DIV2K_train_LR_bicubic
│ │ ├── X2
│ │ ├── X3
│ │ ├── X4
│ │ ├── X2_sub
│ │ ├── X3_sub
│ │ ├── X4_sub
│ ├── DIV2K_train_LR_bicubic_X2_sub.lmdb
│ ├── DIV2K_train_LR_bicubic_X3_sub.lmdb
│ ├── DIV2K_train_LR_bicubic_X4_sub.lmdb
│ ├── ...
这时候有人就会问为什么valid没有执行这一步骤, 因为默认确实不执行, 不过也好说。train_pipeline用lmdb, val_pipeline就用folder呗。毕竟val也不做裁剪。
第三步修改配置文件
如果你之前接触过类似于paddlepaddle, mmcv这类多任务的开源框架的话, 改配置对你来说就肯定不算难事情了。
你如果不改网络结构的话, 也就Augmentation / lr_config / loss and optimizer 和日志以及保存部分需要改。
不过既然是超分, 那就难免有不同的地方。所以就列举几个比较特殊的需要更改的点,
大体例子参照官方提供的EDSR模型配置文件。不是很了解的也可以顺便看看对应参数的解释。
请注意, 该配置文件仅训练2倍下采样的数据。可在scale处修改。
如果使用的是 SRAnnotationDataset 数据集格式进行训练, 则需要修改以下位置:
# data -> train -> dataset
lq_folder='data/DIV2K/DIV2K_train_LR_bicubic/X2_sub'
gt_folder='data/DIV2K/DIV2K_train_HR_sub'
ann_file='data/DIV2K/DIV2K_train_HR_sub.lmdb/meta_info.txt' <--- 经preprocess_div2k_dataset.py处理后生成的txt文件
# data -> val or test -> dataset
lq_folder='./data/val_set5/Set5_bicLRx2'
gt_folder='./data/val_set5/Set5_mod12' <--- 这里可以改为你想要当作验证和测试的数据文件夹, 因为默认是 SRFolderDataset 因此没有ann.txt
如果使用的是 SRLmdbDataset 数据集格式进行训练, 则需要修改以下位置:
# train_dataset_type
train_dataset_type = 'SRLmdbDataset'
# data -> train -> dataset
lq_folder='data/DIV2K/DIV2K_train_LR_bicubic_X2_sub.lmdb'
gt_folder='data/DIV2K/DIV2K_train_HR_sub.lmdb
# 并注释掉原来的ann_file
# train_pipeline, 前两个dict
dict(
type='LoadImageFromFile',
io_backend='lmdb',
key='lq',
db_path='data/DIV2K/DIV2K_train_LR_bicubic_X2_sub.lmdb',
flag='unchanged'),
dict(
type='LoadImageFromFile',
io_backend='lmdb',
key='gt',
db_path='data/DIV2K/DIV2K_train_HR_sub.lmdb',
flag='unchanged')
其他修改:
# evaluation
# 删除gpu_collect项
第四步训练, 懂的都懂。
第五步, 直接用api作推理。
如果你曾经使用过mmcv任何工具框架的api, 可能会像下面这样做:
from mmedit.apis import init_model, restoration_inference
def infer_restoration(img, cfg, ckpt):
model = init_model(config=cfg, checkpoint=ckpt, device="cuda:0")
infer = restoration_inference(model=model, img=item))
return infer
if __name__ == '__main__':
res = infer_restoration(img="xxx.jpg",
cfg="xxx.py",
ckpt="xxx.pth")
然后你会发现, res矩阵里面全都是在0-1之间的数, 你可能会想到是不是做了归一化。
之后一通 *255 再 reshape 再 tensor.numpy(), 发现出来的结果怎么是原低分辨率的图复制了几次而已。
没辙了。
其实有一个叫做tensor2img的函数是需要一起配套使用的, 阅读源码发现是需要再另外做一系列后处理的,
即如下面所示即可:
from mmedit.apis import init_model, restoration_inference
from mmedit.core import tensor2img
def infer_restoration(img, cfg, ckpt):
model = init_model(config=cfg, checkpoint=ckpt, device="cuda:0")
infer = tensor2img(restoration_inference(model=model, img=item)))
return infer
if __name__ == '__main__':
res = infer_restoration(img="xxx.jpg",
cfg="xxx.py",
ckpt="xxx.pth")
