Hu commited on
Commit
33ac1eb
1 Parent(s): 2f110b2

change app name

Browse files
Files changed (1) hide show
  1. app.py +54 -0
app.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import numpy as np
2
+ import gradio as gr
3
+ import torch
4
+ import torch.nn as nn
5
+ import torch.nn.functional as F
6
+ from model import SRCNNModel, pred_SRCNN
7
+ from PIL import Image
8
+
9
+
10
+ title = "Super Resolution with CNN"
11
+ description = """
12
+
13
+ Your low resolution image will be reconstructed to high resolution with a scale of 2 with a convolutional neural network!
14
+
15
+ CNN output on the left, bicubic interpolation output on the right.
16
+
17
+
18
+ """
19
+
20
+ article = "Check out the origianl [paper](https://arxiv.org/abs/1501.00092) proposed by Dong *et al*."
21
+
22
+ # load model
23
+ print("Loading SRCNN model...")
24
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
25
+
26
+ model = SRCNNModel().to(device)
27
+ model.load_state_dict(torch.load('SRCNNmodel_trained.pt'))
28
+ model.eval()
29
+ print("SRCNN model loaded!")
30
+
31
+ def image_grid(imgs, rows, cols):
32
+ '''
33
+ imgs:list of PILImage
34
+ '''
35
+ assert len(imgs) == rows*cols
36
+
37
+ w, h = imgs[0].size
38
+ grid = Image.new('RGB', size=(cols*w, rows*h))
39
+ grid_w, grid_h = grid.size
40
+
41
+ for i, img in enumerate(imgs):
42
+ grid.paste(img, box=(i%cols*w, i//cols*h))
43
+ return grid
44
+
45
+ def sepia(image_path):
46
+ # gradio open image as np array
47
+ image = Image.fromarray(image_path,mode='RGB')
48
+ out_final,image_bicubic,image = pred_SRCNN(model=model,image=image,device=device)
49
+ grid = image_grid([out_final,image_bicubic],1,2)
50
+ return grid
51
+
52
+ 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'])
53
+
54
+ demo.launch(share=True)