Hu commited on
Commit
e6d22e1
·
2 Parent(s): 1ed8271 9f82f91

Merge branch 'main' of https://huggingface.co/spaces/SuwoE/SuperResolution

Browse files
Files changed (1) hide show
  1. demo.py +0 -54
demo.py DELETED
@@ -1,54 +0,0 @@
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)