HuSusu/SuperResolution

1import numpy as np2import gradio as gr3import torch4import torch.nn as nn5import torch.nn.functional as F6import torchvision7from torchvision import transforms8from PIL import Image9 10 11title = "Super Resolution with CNN"12description = """13 14Your low resolution image will be reconstructed to high resolution with a scale of 2 with a convolutional neural network!<br>15 16Detailed training and dataset can be found on my [github repo](https://github.com/susuhu/super-resolution).<br>17 18"""19 20article = """21<div style='margin:20px auto;'>22<p>Sources:<p>23<p>📜 <a href="https://arxiv.org/abs/1501.00092">Image Super-Resolution Using Deep Convolutional Networks</a></p>24<p>📦 Dataset <a href="https://github.com/eugenesiow/super-image-data">this GitHub repo</a></p>25</div>26"""27examples = [28    ["peperoni.png"],29    ["barbara.png"],30]31 32 33class SRCNNModel(nn.Module):34    def __init__(self):35        super(SRCNNModel, self).__init__()36        self.conv1 = nn.Conv2d(1, 64, 9, padding=4)37        self.conv2 = nn.Conv2d(64, 32, 1, padding=0)38        self.conv3 = nn.Conv2d(32, 1, 5, padding=2)39 40    def forward(self, x):41        out = F.relu(self.conv1(x))42        out = F.relu(self.conv2(out))43        out = self.conv3(out)44        return out45 46 47def pred_SRCNN(model, image, device, scale_factor=2):48    """49    model: SRCNN model50    image: low resolution image PILLOW image51    scale_factor: scale factor for resolution52    device: cuda or cpu53    """54    model.to(device)55    model.eval()56 57    # open image, gradio opens image as nparray58    image = Image.fromarray(image)59    # split channels60    y, cb, cr = image.convert("YCbCr").split()61    # size will be used in image transform62    original_size = y.size63 64    # bicubic interpolate it to the original size65    y_bicubic = transforms.Resize(66        (original_size[1] * scale_factor, original_size[0] * scale_factor),67        interpolation=transforms.InterpolationMode.BICUBIC,68    )(y)69    cb_bicubic = transforms.Resize(70        (original_size[1] * scale_factor, original_size[0] * scale_factor),71        interpolation=transforms.InterpolationMode.BICUBIC,72    )(cb)73    cr_bicubic = transforms.Resize(74        (original_size[1] * scale_factor, original_size[0] * scale_factor),75        interpolation=transforms.InterpolationMode.BICUBIC,76    )(cr)77    # turn it into tensor and add batch dimension78    y_bicubic = transforms.ToTensor()(y_bicubic).to(device).unsqueeze(0)79    # get the y channel SRCNN prediction80    y_pred = model(y_bicubic)81    # convert it to numpy image82    y_pred = y_pred[0].cpu().detach().numpy()83 84    # convert it into regular image pixel values85    y_pred = y_pred * 25586    y_pred.clip(0, 255)87    # conver y channel from array to PIL image format for merging88    y_pred_PIL = Image.fromarray(np.uint8(y_pred[0]), mode="L")89    # merge the SRCNN y channel with cb cr channels90    out_final = Image.merge("YCbCr", [y_pred_PIL, cb_bicubic, cr_bicubic]).convert(91        "RGB"92    )93 94    image_bicubic = transforms.Resize(95        (original_size[1] * scale_factor, original_size[0] * scale_factor),96        interpolation=transforms.InterpolationMode.BICUBIC,97    )(image)98    return out_final, image_bicubic99 100 101# load model102# print("Loading  SRCNN model...")103device = torch.device("cuda" if torch.cuda.is_available() else "cpu")104 105model = SRCNNModel().to(device)106model.load_state_dict(107    torch.load("SRCNNmodel_trained.pt", map_location=torch.device(device))108)109model.eval()110# print("SRCNN model loaded!")111 112 113# def image_grid(imgs, rows, cols):114#     '''115#     imgs:list of PILImage116#     '''117#     assert len(imgs) == rows*cols118 119#     w, h = imgs[0].size120#     grid = Image.new('RGB', size=(cols*w, rows*h))121#     grid_w, grid_h = grid.size122 123#     for i, img in enumerate(imgs):124#         grid.paste(img, box=(i%cols*w, i//cols*h))125#     return grid126 127 128def super_reso(input_image):129    # gradio open image as np array130    #image_array = np.asarray(image_path)131    #image = Image.fromarray(image_array, mode="RGB") 132 133    # prediction134    with torch.no_grad():135        out_final, image_bicubic = pred_SRCNN(136            model=model, image=input_image, device=device137        )138    # grid = image_grid([out_final,image_bicubic],1,2)139    return out_final, image_bicubic140 141 142gr.Interface(143    fn=super_reso,144    inputs=gr.Image(label="Upload image"),145    outputs=[146        gr.Image(label="Convolutional neural network"),147        gr.Image(label="Bicubic interpoloation"),148    ],149    title=title,150    description=description,151    article=article,152    examples=examples,153).launch()154 155 156# TypeError: AsyncConnectionPool.__init__() got an unexpected keyword argument 'socket_options'