-
Notifications
You must be signed in to change notification settings - Fork 0
/
Super_Resolution.py
85 lines (80 loc) · 2.79 KB
/
Super_Resolution.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
import cv2
def super_res(img, scale, model="edsr"):
"""
img: Input image array
models: edsr(scale: 2, 3, 4), espcn(scale: 2, 3, 4), fsrcnn(scale: 2, 3, 4), lapsrn(scale: 2, 4, 8)
"""
sr = cv2.dnn_superres.DnnSuperResImpl_create()
if model == "edsr":
if scale == 2:
path = "SuperResModels\EDSR_x2.pb"
sr.readModel(path)
sr.setModel("edsr",2)
elif scale == 3:
path = "SuperResModels\EDSR_x3.pb"
sr.readModel(path)
sr.setModel("edsr",3)
elif scale == 4:
path = "SuperResModels\EDSR_x4.pb"
sr.readModel(path)
sr.setModel("edsr",4)
else:
print("EDSR models only scale at 2x, 3x, and 4x")
elif model == "espcn":
if scale == 2:
path = "SuperResModels\ESPCN_x2.pb"
sr.readModel(path)
sr.setModel("espcn",2)
elif scale == 3:
path = "SuperResModels\ESPCN_x3.pb"
sr.readModel(path)
sr.setModel("espcn",3)
elif scale == 4:
path = "SuperResModels\ESPCN_x4.pb"
sr.readModel(path)
sr.setModel("espcn",4)
else:
print("ESPCN models only scale at 2x, 3x, and 4x")
elif model == "fsrcnn":
if scale == 2:
path = "SuperResModels\FSRCNN_x2.pb"
sr.readModel(path)
sr.setModel("fsrcnn",2)
elif scale == 3:
path = "SuperResModels\FSRCNN_x3.pb"
sr.readModel(path)
sr.setModel("fsrcnn",3)
elif scale == 4:
path = "SuperResModels\FSRCNN_x4.pb"
sr.readModel(path)
sr.setModel("fsrcnn",4)
else:
print("FSRCNN models only scale at 2x, 3x, and 4x")
elif model == "lapsrn":
if scale == 2:
path = "SuperResModels\LapSRN_x2.pb"
sr.readModel(path)
sr.setModel("lapsrn",2)
elif scale == 4:
path = "SuperResModels\LapSRN_x4.pb"
sr.readModel(path)
sr.setModel("lapsrn",4)
elif scale == 8:
path = "SuperResModels\LapSRN_x8.pb"
sr.readModel(path)
sr.setModel("lapsrn",8)
else:
print("LapSRN models only scale at 2x, 4x, and 8x")
result = sr.upsample(img)
return result[:,:,::-1]
if __name__ == "__main__":
# Original image
img = cv2.imread(r"C:\Users\Dennis Pkemoi\Desktop\Coding\PythonProjects\Vidops\Examples\deblurred70.jpg")
cv2.imshow("original Image", img)
# Upscale
model = "fsrcnn"
scale = 2
res = cv2.cvtColor(super_res(img, scale=scale, model=model), cv2.COLOR_BGR2RGB)
cv2.imshow("Super Resolution", res)
cv2.imwrite("Examples\{}_{}.jpg".format(model, str(scale)), res)
cv2.waitKey(0)