BSRGAN的实现

122次阅读
没有评论

共计 3795 个字符,预计需要花费 10 分钟才能阅读完成。

提醒:本文最后更新于 2024-08-30 15:38,文中所关联的信息可能已发生改变,请知悉!

测试代码

import os.path
import logging
import torch

from utils import utils_logger
from utils import utils_image as util
from models.network_rrdbnet import RRDBNet as net

def main():

    utils_logger.logger_info('blind_sr_log', log_path='blind_sr_log.log')
    logger = logging.getLogger('blind_sr_log')

    testsets = 'testsets'       # fixed, set path of testsets
    testset_Ls = ['RealSRSet']  # ['RealSRSet','DPED']

    model_names = ['RRDB','ESRGAN','FSSR_DPED','FSSR_JPEG','RealSR_DPED','RealSR_JPEG']
    model_names = ['BSRGAN']    # 'BSRGANx2' for scale factor 2

    save_results = True
    sf = 4
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    for model_name in model_names:
        if model_name in ['BSRGANx2']:
            sf = 2
        model_path = os.path.join('model_zoo', model_name+'.pth') # set model path
        logger.info('{:>16s} : {:s}'.format('Model Name', model_name))

        # torch.cuda.set_device(0)      # set GPU ID
        logger.info('{:>16s} : {:<d}'.format('GPU ID', torch.cuda.current_device()))
        torch.cuda.empty_cache()

        # --------------------------------
        # define network and load model
        # --------------------------------
        model = net(in_nc=3, out_nc=3, nf=64, nb=23, gc=32, sf=sf)  # define network

        model.load_state_dict(torch.load(model_path), strict=True)
        model.eval()
        for k, v in model.named_parameters():
            v.requires_grad = False
        model = model.to(device)
        torch.cuda.empty_cache()

        for testset_L in testset_Ls:

            L_path = os.path.join(testsets, testset_L)
            #E_path = os.path.join(testsets, testset_L+'_'+model_name)
            E_path = os.path.join(testsets, testset_L+'_results_x'+str(sf))
            util.mkdir(E_path)

            logger.info('{:>16s} : {:s}'.format('Input Path', L_path))
            logger.info('{:>16s} : {:s}'.format('Output Path', E_path))
            idx = 0

            for img in util.get_image_paths(L_path):

                # --------------------------------
                # (1) img_L
                # --------------------------------
                idx += 1
                img_name, ext = os.path.splitext(os.path.basename(img))
                logger.info('{:->4d} --> {:<s} --> x{:<d}--> {:<s}'.format(idx, model_name, sf, img_name+ext))

                img_L = util.imread_uint(img, n_channels=3)
                img_L = util.uint2tensor4(img_L)
                img_L = img_L.to(device)

                # --------------------------------
                # (2) inference
                # --------------------------------
                img_E = model(img_L)

                # --------------------------------
                # (3) img_E
                # --------------------------------
                img_E = util.tensor2uint(img_E)
                if save_results:
                    util.imsave(img_E, os.path.join(E_path, img_name+'_'+model_name+'.png'))

if __name__ == '__main__':

    main()

结果示例

LogHandlers setup!
22-07-26 21:04:13.379 :       Model Name : BSRGAN
22-07-26 21:04:13.386 :           GPU ID : 0
[3, 3, 64, 23, 32, 4]
22-07-26 21:04:14.677 :       Input Path : testsets\RealSRSet
22-07-26 21:04:14.678 :      Output Path : testsets\RealSRSet_results_x4
22-07-26 21:04:14.679 : ---1 --> BSRGAN --> x4--> Lincoln.png
22-07-26 21:04:20.763 : ---2 --> BSRGAN --> x4--> building.png
22-07-26 21:04:21.444 : ---3 --> BSRGAN --> x4--> butterfly.png
22-07-26 21:04:23.191 : ---4 --> BSRGAN --> x4--> butterfly2.png
22-07-26 21:04:23.802 : ---5 --> BSRGAN --> x4--> chip.png
22-07-26 21:04:24.154 : ---6 --> BSRGAN --> x4--> comic1.png
22-07-26 21:04:24.544 : ---7 --> BSRGAN --> x4--> comic2.png
22-07-26 21:04:25.381 : ---8 --> BSRGAN --> x4--> comic3.png
22-07-26 21:04:27.064 : ---9 --> BSRGAN --> x4--> computer.png
22-07-26 21:04:28.882 : --10 --> BSRGAN --> x4--> dog.png
22-07-26 21:04:31.538 : --11 --> BSRGAN --> x4--> dped_crop00061.png
22-07-26 21:04:33.001 : --12 --> BSRGAN --> x4--> foreman.png
22-07-26 21:04:34.198 : --13 --> BSRGAN --> x4--> frog.png
22-07-26 21:04:35.514 : --14 --> BSRGAN --> x4--> oldphoto2.png
22-07-26 21:04:36.157 : --15 --> BSRGAN --> x4--> oldphoto3.png
22-07-26 21:04:36.722 : --16 --> BSRGAN --> x4--> oldphoto6.png
22-07-26 21:04:38.394 : --17 --> BSRGAN --> x4--> painting.png
22-07-26 21:04:39.772 : --18 --> BSRGAN --> x4--> pattern.png
22-07-26 21:04:40.500 : --19 --> BSRGAN --> x4--> ppt3.png
22-07-26 21:04:41.786 : --20 --> BSRGAN --> x4--> tiger.png

部分图片

BSRGAN 的实现

BSRGAN 的实现

BSRGAN 的实现

BSRGAN 的实现

正文完
 0
icvuln
版权声明:本站原创文章,由 icvuln 于2022-07-24发表,共计3795字。
转载说明:除特殊说明外本站文章皆由CC-4.0协议发布,转载请注明出处。
评论(没有评论)