基于Python制作图像完美超分处理工具_Python

前言

根据上一篇文章：Python图片视频超分模型RealBasicVSR的使用教程

我把项目进行的工具化调整，做成了对单一照片输入进行处理的工具。大家可以拿去直接集成到代码中使用。

github项目地址：工具项目地址

项目说明

项目结构如下

基于Python制作图像完美超分处理工具

注意项

1、 checkpoints中的模型文件太大，上传不到github上，大家可以通过下面我分享的百度云盘地址下载，放到checkpoints目录中。

链接提取码：1234

2、conda环境安装，我这里就不多叙述了，直接参考上一篇文章中的安装过程即可。

Python图片视频超分模型RealBasicVSR的使用教程

代码说明

主要的核心代码是inference_hy.py文件，代码如下：

				?

									#!/user/bin/env python

									# coding=utf-8

									"""

									@project : hy_realbasicvsr

									@author  : 剑客阿良_ALiang

									@file   : inference_hy.py

									@ide    : PyCharm

									@time   : 2022-05-08 15:18:15

									"""

									import argparse

									import glob

									import os

									import cv2

									import mmcv

									import numpy as np

									import torch

									import uuid

									from mmcv.runner import load_checkpoint

									from mmedit.core import tensor2img

									from realbasicvsr.models.builder import build_model

									def init_model(config, checkpoint=None):

									    if isinstance(config, str):

									        config = mmcv.Config.fromfile(config)

									    elif not isinstance(config, mmcv.Config):

									        raise TypeError('config must be a filename or Config object, '

									                        f'but got {type(config)}')

									    config.model.pretrained = None

									    config.test_cfg.metrics = None

									    model = build_model(config.model, test_cfg=config.test_cfg)

									    if checkpoint is not None:

									        checkpoint = load_checkpoint(model, checkpoint)

									    model.cfg = config  # save the config in the model for convenience

									    model.eval()

									    return model

									class Worker:

									    def __init__(self):

									        self.checkpoint_path = 'checkpoints/RealBasicVSR_x4.pth'

									        self.config = 'configs/realbasicvsr_x4.py'

									        self.is_save_as_png = True

									        self.max_seq_len = 2

									        self.model = init_model(self.config, self.checkpoint_path)

									    def do_pic(self, input_image_path: str, output_dir: str):

									        inputs = []

									        img = mmcv.imread(input_image_path, channel_order='rgb')

									        ext = os.path.basename(input_image_path).split('.')[-1]

									        inputs.append(img)

									        for i, img in enumerate(inputs):

									            img = torch.from_numpy(img / 255.).permute(2, 0, 1).float()

									            inputs[i] = img.unsqueeze(0)

									        inputs = torch.stack(inputs, dim=1)

									        # map to cuda, if available

									        cuda_flag = False

									        if torch.cuda.is_available():

									            model = self.model.cuda()

									            cuda_flag = True

									        with torch.no_grad():

									            if isinstance(self.max_seq_len, int):

									                outputs = []

									                for i in range(0, inputs.size(1), self.max_seq_len):

									                    imgs = inputs[:, i:i + self.max_seq_len, :, :, :]

									                    if cuda_flag:

									                        imgs = imgs.cuda()

									                    outputs.append(self.model(imgs, test_mode=True)['output'].cpu())

									                outputs = torch.cat(outputs, dim=1)

									            else:

									                if cuda_flag:

									                    inputs = inputs.cuda()

									                outputs = self.model(inputs, test_mode=True)['output'].cpu()

									        mmcv.mkdir_or_exist(output_dir)

									        for i in range(0, outputs.size(1)):

									            output = tensor2img(outputs[:, i, :, :, :])

									            filename = '{}.{}'.format(uuid.uuid1().hex, ext)

									            if self.is_save_as_png:

									                file_extension = os.path.splitext(filename)[1]

									                filename = filename.replace(file_extension, '.png')

									            result_path = os.path.join(output_dir, filename)

									            mmcv.imwrite(output, result_path)

									if __name__ == '__main__':

									    worker = Worker()

									    worker.do_pic('data/136.jpeg', 'results/')

代码说明

1、如果在真实项目中使用的话，worker对象可以在项目启动中创建，不需要每次调用时候创建。

2、主要使用方法为do_pic方法，入参分别为：需要超分的照片地址，以及最后生成的目录地址。

3、最终输出的照片名使用uuid进行命名，如果需要自定义的话，可以看着调整。

4、输出的文件格式不一定需要是png，上面的类初始化有bool选项，自行调整。

测试一下代码，先看一下图片。

基于Python制作图像完美超分处理工具

执行结果

2022-05-08 15:47:54,792 - mmedit - INFO - load checkpoint from torchvision path: torchvision://vgg19
load checkpoint from local path: checkpoints/RealBasicVSR_x4.pth
/home/zhangmaolin/.local/lib/python3.7/site-packages/torch/functional.py:445: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at ../aten/src/ATen/native/TensorShape.cpp:2157.)
return _VF.meshgrid(tensors, **kwargs) # type: ignore[attr-defined]

最终图片输出

基于Python制作图像完美超分处理工具