import cv2
import numpy as np
import torch

from basicsr.archs.rrdbnet_arch import RRDBNet


def init_sr_model(model_path):
    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32)
    model.load_state_dict(torch.load(model_path)['params'], strict=True)
    model.eval()
    model = model.cuda()
    return model


def enhance(model, image):
    img = image.astype(np.float32) / 255.
    img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
    img = img.unsqueeze(0).cuda()
    with torch.no_grad():
        output = model(img)
    output = output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
    output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
    output = (output * 255.0).round().astype(np.uint8)
    return output