initial commit

SRCNNmodel_trained.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c1c33c257abf0eef36eb73c60fbc1863ebf7612cefb07c6a7aea85b283b03ddb
+size 34455

demo.py

@@ -0,0 +1,54 @@
+import numpy as np
+import gradio as gr
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from model import SRCNNModel, pred_SRCNN
+from PIL import Image
+
+
+title = "Super Resolution with CNN"
+description = """
+
+Your low resolution image will be reconstructed to high resolution with a scale of 2 with a convolutional neural network!
+
+CNN output on the left, bicubic interpolation output on the right.
+
+
+"""
+
+article = "Check out the origianl [paper](https://arxiv.org/abs/1501.00092) proposed by Dong *et al*."
+
+# load model
+print("Loading  SRCNN model...")
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+model = SRCNNModel().to(device)
+model.load_state_dict(torch.load('SRCNNmodel_trained.pt'))
+model.eval()
+print("SRCNN model loaded!")
+
+def image_grid(imgs, rows, cols):
+    '''
+    imgs:list of PILImage
+    '''
+    assert len(imgs) == rows*cols
+
+    w, h = imgs[0].size
+    grid = Image.new('RGB', size=(cols*w, rows*h))
+    grid_w, grid_h = grid.size
+    
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i%cols*w, i//cols*h))
+    return grid
+
+def sepia(image_path):
+    # gradio open image as np array
+    image = Image.fromarray(image_path,mode='RGB')
+    out_final,image_bicubic,image = pred_SRCNN(model=model,image=image,device=device)
+    grid = image_grid([out_final,image_bicubic],1,2)
+    return grid
+
+demo = gr.Interface(fn = sepia, inputs=gr.Image(shape=(200, 200)), outputs="image",title=title,description = description,article = article,examples=['LR_image.png','barbara.png'])
+
+demo.launch(share=True)

model.py

@@ -0,0 +1,80 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+from torchvision.transforms import transforms
+import numpy as np
+from PIL import Image
+
+class SRCNNModel(nn.Module):
+  def __init__(self):
+    super(SRCNNModel, self).__init__()
+    self.conv1=nn.Conv2d(1,64,9,padding=4)
+    self.conv2=nn.Conv2d(64,32,1,padding=0)
+    self.conv3=nn.Conv2d(32,1,5,padding=2)
+
+  def forward(self,x):
+    out = F.relu(self.conv1(x))
+    out = F.relu(self.conv2(out))
+    out = self.conv3(out)
+    return out
+
+def pred_SRCNN(model,image,device,scale_factor=2):
+    """
+    model: SRCNN model
+    image: low resolution image PILLOW image
+    scale_factor: scale factor for resolution
+    device: cuda or cpu
+    """
+    model.to(device)
+    model.eval()
+
+    # open image
+    # image = Image.open(image_path)
+    # split channels
+    y, cb, cr= image.convert('YCbCr').split()
+    # size will be used in image transform
+    original_size = y.size 
+
+    # bicubic interpolate it to the original size
+    y_bicubic = transforms.Resize((original_size[1]*scale_factor,original_size[0]*scale_factor),interpolation=Image.BICUBIC)(y)
+    cb_bicubic = transforms.Resize((original_size[1]*scale_factor,original_size[0]*scale_factor),interpolation=Image.BICUBIC)(cb)
+    cr_bicubic = transforms.Resize((original_size[1]*scale_factor,original_size[0]*scale_factor),interpolation=Image.BICUBIC)(cr)
+    # turn it into tensor and add batch dimension
+    y_bicubic = transforms.ToTensor()(y_bicubic).to(device).unsqueeze(0) 
+    # get the y channel SRCNN prediction
+    y_pred = model(y_bicubic)
+    # convert it to numpy image
+    y_pred = y_pred[0].cpu().detach().numpy()
+
+    # convert it into regular image pixel values
+    y_pred = y_pred*255 
+    y_pred.clip(0,255)
+    # conver y channel from array to PIL image format for merging
+    y_pred_PIL = Image.fromarray(np.uint8(y_pred[0]),mode='L')
+    # merge the SRCNN y channel with cb cr channels
+    out_final = Image.merge('YCbCr',[y_pred_PIL,cb_bicubic,cr_bicubic]).convert('RGB') 
+
+    image_bicubic = transforms.Resize((original_size[1]*scale_factor,original_size[0]*scale_factor),interpolation=Image.BICUBIC)(image)
+    return out_final,image_bicubic,image
+
+
+def main():
+    print("Loading  SRCNN model...")
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    model = SRCNNModel().to(device)
+    model.load_state_dict(torch.load('SRCNNmodel_trained.pt'))
+    model.eval()
+    print("SRCNN model loaded!")
+
+    image_path = "LR_image.png"
+
+    out_final,image_bicubic,image = pred_SRCNN(model=model,image_path=image_path,device=device)
+    image.show()
+    out_final.show()
+    image_bicubic.show()
+
+
+if __name__=="__main__":
+    main()