Inference issue about ecbsr

[puppy9207]

hi, can you share your test code? i want to convert ychannel got by ECBSR and CbCr channel got by bicubic to rgb, but the final image contains some odd pixels

Here is my code. And there were two problems above. The first was that the file type was changed to a changed jpg extension, and the second was that the pre-trained file was not properly retrieved. It was my mistake.

import torch
import skimage.color as sc
import cv2
import glob
import os
import time
import argparse, yaml
import numpy as np
from models.ecbsr import ECBSR
from models.plainsr import PlainSR
""" python infer.py --config experiments/ecbsr-x2-m4c16-prelu-2021-1008-0705/config.yml \
--pretrain experiments/ecbsr-x2-m4c16-prelu-2021-1008-0705/models/model_x2_999.pt"""

parser = argparse.ArgumentParser(description='ECBSR')

parser.add_argument('--config', type=str, default=None, help = 'pre-config file for training')
parser.add_argument('--scale', type=int, default=2, help = 'scale for sr network')
parser.add_argument('--colors', type=int, default=1, help = '1(Y channls of YCbCr)')
parser.add_argument('--m_ecbsr', type=int, default=4, help = 'number of ecb')
parser.add_argument('--c_ecbsr', type=int, default=16, help = 'channels of ecb')
parser.add_argument('--idt_ecbsr', type=int, default=0, help = 'incorporate identity mapping in ecb or not')
parser.add_argument('--act_type', type=str, default='prelu', help = 'prelu, relu, splus, rrelu')
parser.add_argument('--pretrain', type=str, default=None, help = 'path of pretrained model')
parser.add_argument('--img_path', type=str, default="imgs/input", help = 'path of images')
parser.add_argument('--save_path', type=str, default="imgs/output", help = 'path of save images')
parser.add_argument('--gpu_id', type=int, default=0, help = 'gpu id for training')
def convert_rgb_to_y(img, dim_order='hwc'):
    if dim_order == 'hwc':
        return 16. + (64.738 * img[..., 0] + 129.057 * img[..., 1] + 25.064 * img[..., 2]) / 256.
    else:
        return 16. + (64.738 * img[0] + 129.057 * img[1] + 25.064 * img[2]) / 256.


def convert_ycbcr_to_rgb(img, dim_order='hwc'):
    if dim_order == 'hwc':
        r = 298.082 * img[..., 0] / 256. + 408.583 * img[..., 2] / 256. - 222.921
        g = 298.082 * img[..., 0] / 256. - 100.291 * img[..., 1] / 256. - 208.120 * img[..., 2] / 256. + 135.576
        b = 298.082 * img[..., 0] / 256. + 516.412 * img[..., 1] / 256. - 276.836
    else:
        r = 298.082 * img[0] / 256. + 408.583 * img[2] / 256. - 222.921
        g = 298.082 * img[0] / 256. - 100.291 * img[1] / 256. - 208.120 * img[2] / 256. + 135.576
        b = 298.082 * img[0] / 256. + 516.412 * img[1] / 256. - 276.836
    return np.array([r, g, b]).transpose([1, 2, 0])

def convert_rgb_to_ycbcr(img, dim_order='hwc'):
    if dim_order == 'hwc':
        y = 16. + (64.738 * img[..., 0] + 129.057 * img[..., 1] + 25.064 * img[..., 2]) / 256.
        cb = 128. + (-37.945 * img[..., 0] - 74.494 * img[..., 1] + 112.439 * img[..., 2]) / 256.
        cr = 128. + (112.439 * img[..., 0] - 94.154 * img[..., 1] - 18.285 * img[..., 2]) / 256.
    else:
        y = 16. + (64.738 * img[0] + 129.057 * img[1] + 25.064 * img[2]) / 256.
        cb = 128. + (-37.945 * img[0] - 74.494 * img[1] + 112.439 * img[2]) / 256.
        cr = 128. + (112.439 * img[0] - 94.154 * img[1] - 18.285 * img[2]) / 256.
    return np.array([y, cb, cr]).transpose([1, 2, 0])

def preprocess(img, device):
    img = np.array(img).astype(np.float32)
    ycbcr = convert_rgb_to_ycbcr(img)
    x = ycbcr[..., 0]
    x = torch.from_numpy(x).to(device)
    x = x.unsqueeze(0).unsqueeze(0)
    return x, ycbcr

if __name__ == '__main__':
    args = parser.parse_args()
    # if args.config:
    #    opt = vars(args)
    #    yaml_args = yaml.load(opearn(gs.config), Loader=yaml.FullLoader)
    #    opt.update(yaml_args)
    
    device = None
    if args.gpu_id >= 0 and torch.cuda.is_available():
        print("use cuda & cudnn for acceleration!")
        print("the gpu id is: {}".format(args.gpu_id))
        device = torch.device('cuda:{}'.format(args.gpu_id))
        torch.backends.cudnn.benchmark = True

    model = ECBSR(module_nums=4, channel_nums=16, with_idt=0, act_type="prelu", scale=2, colors=1).to(device)
    model_plain = PlainSR(module_nums=4, channel_nums=16, act_type="prelu", scale=2, colors=1).to(device)
    print("load pretrained model: {}!".format(args.pretrain))
    model.load_state_dict(torch.load(args.pretrain,map_location=device))

    depth = len(model.backbone)
    for d in range(depth):
        module = model.backbone[d]
        act_type = module.act_type
        RK, RB = module.rep_params()
        model_plain.backbone[d].conv3x3.weight.data = RK
        model_plain.backbone[d].conv3x3.bias.data = RB

        if act_type == 'relu':     pass
        elif act_type == 'linear': pass
        elif act_type == 'prelu':  model_plain.backbone[d].act.weight.data = module.act.weight.data
        else: raise ValueError('invalid type of activation!')

    img_list = glob.glob(args.img_path+'/*')
    save_dir= args.save_path
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    for img_path in img_list:
        img = cv2.imread(img_path)
        start = time.time()
        # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        pre_img, ycbcr = preprocess(img,device)

        ycbcr = cv2.resize(ycbcr,(0,0),fx=2,fy=2)
        # print(img.shape)
        # pre_img = img[:, :, 0:1]
        # print(pre_img.shape)
        # tensor = torch.from_numpy(pre_img).float().unsqueeze(0).to(device)
        
        out = model(pre_img).cpu().detach().numpy().squeeze(0).squeeze(0)
        print(out.shape)
        output = np.array([out, ycbcr[..., 1], ycbcr[..., 2]]).transpose([1, 2, 0])
        output = np.clip(convert_ycbcr_to_rgb(output), 0.0, 255.0).astype(np.uint8)
        # output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
        print(time.time()-start)
        cv2.imwrite(os.path.join(save_dir,os.path.basename(img_path)),output)
    # lr_image, hr_image = sc.rgb2ycbcr(lr_image)[:, :, 0:1], sc.rgb2ycbcr(hr_image)[:, :, 0:1]```

[fangichao]

hi, can you share your test code? i want to convert ychannel got by ECBSR and CbCr channel got by bicubic to rgb, but the final image contains some odd pixels

[fangichao]

thks a lot, problem is solved

[fangichao]

hi, can you share your test code? i want to convert ychannel got by ECBSR and CbCr channel got by bicubic to rgb, but the final image contains some odd pixels

Here is my code. And there were two problems above. The first was that the file type was changed to a changed jpg extension, and the second was that the pre-trained file was not properly retrieved. It was my mistake.

import torch
import skimage.color as sc
import cv2
import glob
import os
import time
import argparse, yaml
import numpy as np
from models.ecbsr import ECBSR
from models.plainsr import PlainSR
""" python infer.py --config experiments/ecbsr-x2-m4c16-prelu-2021-1008-0705/config.yml \
--pretrain experiments/ecbsr-x2-m4c16-prelu-2021-1008-0705/models/model_x2_999.pt"""

parser = argparse.ArgumentParser(description='ECBSR')

parser.add_argument('--config', type=str, default=None, help = 'pre-config file for training')
parser.add_argument('--scale', type=int, default=2, help = 'scale for sr network')
parser.add_argument('--colors', type=int, default=1, help = '1(Y channls of YCbCr)')
parser.add_argument('--m_ecbsr', type=int, default=4, help = 'number of ecb')
parser.add_argument('--c_ecbsr', type=int, default=16, help = 'channels of ecb')
parser.add_argument('--idt_ecbsr', type=int, default=0, help = 'incorporate identity mapping in ecb or not')
parser.add_argument('--act_type', type=str, default='prelu', help = 'prelu, relu, splus, rrelu')
parser.add_argument('--pretrain', type=str, default=None, help = 'path of pretrained model')
parser.add_argument('--img_path', type=str, default="imgs/input", help = 'path of images')
parser.add_argument('--save_path', type=str, default="imgs/output", help = 'path of save images')
parser.add_argument('--gpu_id', type=int, default=0, help = 'gpu id for training')
def convert_rgb_to_y(img, dim_order='hwc'):
    if dim_order == 'hwc':
        return 16. + (64.738 * img[..., 0] + 129.057 * img[..., 1] + 25.064 * img[..., 2]) / 256.
    else:
        return 16. + (64.738 * img[0] + 129.057 * img[1] + 25.064 * img[2]) / 256.


def convert_ycbcr_to_rgb(img, dim_order='hwc'):
    if dim_order == 'hwc':
        r = 298.082 * img[..., 0] / 256. + 408.583 * img[..., 2] / 256. - 222.921
        g = 298.082 * img[..., 0] / 256. - 100.291 * img[..., 1] / 256. - 208.120 * img[..., 2] / 256. + 135.576
        b = 298.082 * img[..., 0] / 256. + 516.412 * img[..., 1] / 256. - 276.836
    else:
        r = 298.082 * img[0] / 256. + 408.583 * img[2] / 256. - 222.921
        g = 298.082 * img[0] / 256. - 100.291 * img[1] / 256. - 208.120 * img[2] / 256. + 135.576
        b = 298.082 * img[0] / 256. + 516.412 * img[1] / 256. - 276.836
    return np.array([r, g, b]).transpose([1, 2, 0])

def convert_rgb_to_ycbcr(img, dim_order='hwc'):
    if dim_order == 'hwc':
        y = 16. + (64.738 * img[..., 0] + 129.057 * img[..., 1] + 25.064 * img[..., 2]) / 256.
        cb = 128. + (-37.945 * img[..., 0] - 74.494 * img[..., 1] + 112.439 * img[..., 2]) / 256.
        cr = 128. + (112.439 * img[..., 0] - 94.154 * img[..., 1] - 18.285 * img[..., 2]) / 256.
    else:
        y = 16. + (64.738 * img[0] + 129.057 * img[1] + 25.064 * img[2]) / 256.
        cb = 128. + (-37.945 * img[0] - 74.494 * img[1] + 112.439 * img[2]) / 256.
        cr = 128. + (112.439 * img[0] - 94.154 * img[1] - 18.285 * img[2]) / 256.
    return np.array([y, cb, cr]).transpose([1, 2, 0])

def preprocess(img, device):
    img = np.array(img).astype(np.float32)
    ycbcr = convert_rgb_to_ycbcr(img)
    x = ycbcr[..., 0]
    x = torch.from_numpy(x).to(device)
    x = x.unsqueeze(0).unsqueeze(0)
    return x, ycbcr

if __name__ == '__main__':
    args = parser.parse_args()
    # if args.config:
    #    opt = vars(args)
    #    yaml_args = yaml.load(opearn(gs.config), Loader=yaml.FullLoader)
    #    opt.update(yaml_args)
    
    device = None
    if args.gpu_id >= 0 and torch.cuda.is_available():
        print("use cuda & cudnn for acceleration!")
        print("the gpu id is: {}".format(args.gpu_id))
        device = torch.device('cuda:{}'.format(args.gpu_id))
        torch.backends.cudnn.benchmark = True

    model = ECBSR(module_nums=4, channel_nums=16, with_idt=0, act_type="prelu", scale=2, colors=1).to(device)
    model_plain = PlainSR(module_nums=4, channel_nums=16, act_type="prelu", scale=2, colors=1).to(device)
    print("load pretrained model: {}!".format(args.pretrain))
    model.load_state_dict(torch.load(args.pretrain,map_location=device))

    depth = len(model.backbone)
    for d in range(depth):
        module = model.backbone[d]
        act_type = module.act_type
        RK, RB = module.rep_params()
        model_plain.backbone[d].conv3x3.weight.data = RK
        model_plain.backbone[d].conv3x3.bias.data = RB

        if act_type == 'relu':     pass
        elif act_type == 'linear': pass
        elif act_type == 'prelu':  model_plain.backbone[d].act.weight.data = module.act.weight.data
        else: raise ValueError('invalid type of activation!')

    img_list = glob.glob(args.img_path+'/*')
    save_dir= args.save_path
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    for img_path in img_list:
        img = cv2.imread(img_path)
        start = time.time()
        # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        pre_img, ycbcr = preprocess(img,device)

        ycbcr = cv2.resize(ycbcr,(0,0),fx=2,fy=2)
        # print(img.shape)
        # pre_img = img[:, :, 0:1]
        # print(pre_img.shape)
        # tensor = torch.from_numpy(pre_img).float().unsqueeze(0).to(device)
        
        out = model(pre_img).cpu().detach().numpy().squeeze(0).squeeze(0)
        print(out.shape)
        output = np.array([out, ycbcr[..., 1], ycbcr[..., 2]]).transpose([1, 2, 0])
        output = np.clip(convert_ycbcr_to_rgb(output), 0.0, 255.0).astype(np.uint8)
        # output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
        print(time.time()-start)
        cv2.imwrite(os.path.join(save_dir,os.path.basename(img_path)),output)
    # lr_image, hr_image = sc.rgb2ycbcr(lr_image)[:, :, 0:1], sc.rgb2ycbcr(hr_image)[:, :, 0:1]```

thks a lot, problem is solved

[simplew2011]

thanks for share，update onnxruntime etc：

import os
import cv2
import yaml
import glob
import utils
import torch
import numpy as np
from tqdm import tqdm
from models.ecbsr import ECBSR
from models.plainsr import PlainSR
import utils


def convert_rgb_to_y(img, dim_order='hwc'):
    if dim_order == 'hwc':
        return 16. + (64.738 * img[..., 0] + 129.057 * img[..., 1] + 25.064 * img[..., 2]) / 256.
    else:
        return 16. + (64.738 * img[0] + 129.057 * img[1] + 25.064 * img[2]) / 256.


def convert_ycbcr_to_rgb(img, dim_order='hwc'):
    if dim_order == 'hwc':
        r = 298.082 * img[..., 0] / 256. + 408.583 * img[..., 2] / 256. - 222.921
        g = 298.082 * img[..., 0] / 256. - 100.291 * img[..., 1] / 256. - 208.120 * img[..., 2] / 256. + 135.576
        b = 298.082 * img[..., 0] / 256. + 516.412 * img[..., 1] / 256. - 276.836
    else:
        r = 298.082 * img[0] / 256. + 408.583 * img[2] / 256. - 222.921
        g = 298.082 * img[0] / 256. - 100.291 * img[1] / 256. - 208.120 * img[2] / 256. + 135.576
        b = 298.082 * img[0] / 256. + 516.412 * img[1] / 256. - 276.836
    return np.array([r, g, b]).transpose([1, 2, 0])

def convert_rgb_to_ycbcr(img, dim_order='hwc'):
    if dim_order == 'hwc':
        y = 16. + (64.738 * img[..., 0] + 129.057 * img[..., 1] + 25.064 * img[..., 2]) / 256.
        cb = 128. + (-37.945 * img[..., 0] - 74.494 * img[..., 1] + 112.439 * img[..., 2]) / 256.
        cr = 128. + (112.439 * img[..., 0] - 94.154 * img[..., 1] - 18.285 * img[..., 2]) / 256.
    else:
        y = 16. + (64.738 * img[0] + 129.057 * img[1] + 25.064 * img[2]) / 256.
        cb = 128. + (-37.945 * img[0] - 74.494 * img[1] + 112.439 * img[2]) / 256.
        cr = 128. + (112.439 * img[0] - 94.154 * img[1] - 18.285 * img[2]) / 256.
    return np.array([y, cb, cr]).transpose([1, 2, 0])

def preprocess(img, device):
    img = np.array(img).astype(np.float32)
    ycbcr = convert_rgb_to_ycbcr(img)
    x = ycbcr[..., 0]
    x = torch.from_numpy(x).to(device)
    x = x.unsqueeze(0).unsqueeze(0)
    return x, ycbcr



if __name__ == "__main__":

    config_path = "configs/ecbsr_x2_m4c8_prelu.yml"
    args = yaml.load(open(config_path), Loader=yaml.FullLoader)

    device = torch.device('cpu')

    model_ecbsr = ECBSR(module_nums=args["m_ecbsr"], channel_nums=args["c_ecbsr"], with_idt=args["idt_ecbsr"], act_type=args["act_type"], scale=args["scale"], colors=args["colors"]).to(device)
    model_ecbsr.load_state_dict(torch.load(args["pretrain"]))
    model_plain = PlainSR(module_nums=args["m_ecbsr"], channel_nums=args["c_ecbsr"], act_type=args["act_type"], scale=args["scale"], colors=args["colors"]).to(device)

    ## copy weights from ecbsr to plainsr
    depth = len(model_ecbsr.backbone)
    for d in range(depth):
        module = model_ecbsr.backbone[d]
        act_type = module.act_type
        RK, RB = module.rep_params()
        model_plain.backbone[d].conv3x3.weight.data = RK
        model_plain.backbone[d].conv3x3.bias.data = RB

        if act_type == 'relu':     pass
        elif act_type == 'linear': pass
        elif act_type == 'prelu':  model_plain.backbone[d].act.weight.data = module.act.weight.data
        else: raise ValueError('invalid type of activation!')
    model_ecbsr.eval()
    model_plain.eval()

    ###############################################################################
    # from thop import profile
    # from thop import clever_format
    # input = torch.randn(1, 1, 1080, 1920)
    # flops, params = profile(model_ecbsr, inputs=(input,))
    # print("flops(G):", "%.3f" % (flops / 900000000 * 2))
    # flops,params = clever_format([ flops / 900000000 * 2,params], "%.3f")
    # print("params:", params)
    # # model_plain (1, 3, 1080, 1920)
    # # flops(G): 12.460
    # # params: 2.748K

    # # model_ecbsr (1, 3, 1080, 1920)
    # # flops(G): 0.184
    # # params: 40.000B
    
    # input_shape = (1, 1, 1080, 1920)

    # shape_dict = [("input", input_shape)]
    # input_data = torch.randn(input_shape)
    # with torch.no_grad():
    #     scripted_model = torch.jit.trace(model_ecbsr, input_data)#.eval()
    #     scripted_model.save("model_ecbsr-1080_1920.torchscript.pt")

    # with torch.no_grad():
    #     torch.onnx.export(model_ecbsr, input_data, "model_ecbsr-1080_1920.onnx", input_names=["input"], output_names=["output"], opset_version=11,
    #     # dynamic_axes= {
    #     #             "input": {0: 'batch_size', 2 : 'in_height', 3: 'in_width'},
    #     #             "output": {0: 'batch_size', 2: 'out_height', 3:'out_width'}}
    #     )
    # exit(0)
    ###############################################################################


    HR_folder = '/data/storage/temp/DIV2K/DIV2K_train_HR'
    LR_folder = '/data/storage/temp/DIV2K/DIV2K_train_LR_bicubic'

    img_list = glob.glob(LR_folder + '/X2/*.png')
    save_dir= "result"
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    for img_path in tqdm(img_list):
        lr = cv2.imread(img_path)
        lr = cv2.resize(lr, [1920, 1080], interpolation=cv2.INTER_CUBIC)
        # hr
        hr = cv2.imread(img_path.replace("DIV2K_train_LR_bicubic", "DIV2K_train_HR").replace("X2/", "").replace("x2", ""))
        lr = cv2.resize(lr, [1920*2, 1080*2], interpolation=cv2.INTER_CUBIC)

        hr_tensor = np.ascontiguousarray(hr.transpose((2, 0, 1)))
        hr_tensor = torch.from_numpy(hr_tensor).float().unsqueeze(0)

        input, ycbcr = preprocess(lr, device)

        with torch.no_grad():
            # out = model_ecbsr(input)
            out = model_plain(input)

        # onnxruntime
        # import onnxruntime
        # onnxruntime_engine = onnxruntime.InferenceSession("model_plain-1080_1920.onnx")
        # onnxruntime_inputs = {onnxruntime_engine.get_inputs()[0].name: np.array(input)}
        # onnxruntime_outputs = onnxruntime_engine.run(None, onnxruntime_inputs)
        # np.testing.assert_allclose(np.array(out), onnxruntime_outputs[0], rtol=1e-02, atol=1e-02)
        # print("Example: Onnx model has been tested with ONNXRuntime, the result looks good !")

        sr = out.cpu().detach().numpy().squeeze(0).squeeze(0)
        # combine sr_y and lr_cb, lr_cr, make last sr
        ycbcr = cv2.resize(ycbcr, (0,0), fx=2,fy=2)
        sr = np.array([sr, ycbcr[..., 1], ycbcr[..., 2]]).transpose([1, 2, 0])
        sr = np.clip(convert_ycbcr_to_rgb(sr), 0.0, 255.0)
        cv2.imwrite(os.path.join(save_dir,os.path.basename(img_path)), sr.astype(np.uint8))

        sr_tensor = torch.from_numpy(sr.transpose((2, 0, 1))).unsqueeze(0)
        psnr = utils.calc_psnr(sr_tensor, hr_tensor)       
        ssim = utils.calc_ssim(sr_tensor, hr_tensor)
        print("{} ---> psnr: {}, ssim: {}".format(os.path.basename(img_path), psnr, ssim))