CHEONMA010 commited on
Commit
5b600b5
1 Parent(s): 7dcb35e

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +65 -65
app.py CHANGED
@@ -1,65 +1,65 @@
1
- import gradio as gr
2
- from PIL import Image
3
- import torch
4
- import numpy as np
5
- import sys
6
- import os
7
- import RRDBNet_arch as arch
8
- # Add the `scripts/` folder to the system path
9
- sys.path.append(os.path.join(os.path.dirname(__file__), 'scripts'))
10
-
11
-
12
-
13
-
14
- # Path to the pretrained model
15
- model_path = os.path.join(os.path.dirname(__file__), 'models', 'RRDB_ESRGAN_x4.pth')
16
-
17
- # Load ESRGAN model
18
- device = torch.device('cpu') #'cuda' if using GPU
19
- model = arch.RRDBNet(3, 3, 64, 23, gc=32)
20
- model.load_state_dict(torch.load(model_path, map_location=device), strict=True)
21
- model.eval()
22
- model = model.to(device)
23
-
24
-
25
- def upscale_image(image):
26
- img = np.array(image) / 255.0
27
- img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
28
- img_LR = img.unsqueeze(0).to(device)
29
-
30
- with torch.no_grad():
31
- output = model(img_LR).data.squeeze().float().cpu().clamp_(0, 1).numpy()
32
- output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
33
- output = (output * 255.0).round().astype(np.uint8)
34
- return Image.fromarray(output)
35
-
36
- # Gradio interface
37
- def gradio_interface(image):
38
- try:
39
-
40
- if image is None:
41
- raise ValueError("No image uploaded. Please upload an image to upscale.")
42
-
43
-
44
- upscaled_image = upscale_image(image)
45
- original_size = image.size
46
- upscaled_size = upscaled_image.size
47
-
48
- return image, upscaled_image, f"Original Size: {original_size[0]}x{original_size[1]}", f"Upscaled Size: {upscaled_size[0]}x{upscaled_size[1]}"
49
-
50
- except Exception as e:
51
-
52
- return None, None, "Error", str(e)
53
-
54
-
55
- gr_interface = gr.Interface(
56
- fn=gradio_interface,
57
- inputs=gr.Image(type="pil"),
58
- outputs=[gr.Image(), gr.Image(), gr.Text(), gr.Text()],
59
- title="ESRGAN Image Upscaler",
60
- description="Upload an image to upscale it using ESRGAN."
61
- )
62
-
63
- if __name__ == '__main__':
64
-
65
- gr_interface.launch()
 
1
+ import gradio as gr
2
+ from PIL import Image
3
+ import torch
4
+ import numpy as np
5
+ import sys
6
+ import os
7
+
8
+ # Add the `scripts/` folder to the system path
9
+ sys.path.append(os.path.join(os.path.dirname(__file__), 'scripts'))
10
+ import RRDBNet_arch as arch
11
+
12
+
13
+
14
+ # Path to the pretrained model
15
+ model_path = os.path.join(os.path.dirname(__file__), 'models', 'RRDB_ESRGAN_x4.pth')
16
+
17
+ # Load ESRGAN model
18
+ device = torch.device('cpu') #'cuda' if using GPU
19
+ model = arch.RRDBNet(3, 3, 64, 23, gc=32)
20
+ model.load_state_dict(torch.load(model_path, map_location=device), strict=True)
21
+ model.eval()
22
+ model = model.to(device)
23
+
24
+
25
+ def upscale_image(image):
26
+ img = np.array(image) / 255.0
27
+ img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
28
+ img_LR = img.unsqueeze(0).to(device)
29
+
30
+ with torch.no_grad():
31
+ output = model(img_LR).data.squeeze().float().cpu().clamp_(0, 1).numpy()
32
+ output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
33
+ output = (output * 255.0).round().astype(np.uint8)
34
+ return Image.fromarray(output)
35
+
36
+ # Gradio interface
37
+ def gradio_interface(image):
38
+ try:
39
+
40
+ if image is None:
41
+ raise ValueError("No image uploaded. Please upload an image to upscale.")
42
+
43
+
44
+ upscaled_image = upscale_image(image)
45
+ original_size = image.size
46
+ upscaled_size = upscaled_image.size
47
+
48
+ return image, upscaled_image, f"Original Size: {original_size[0]}x{original_size[1]}", f"Upscaled Size: {upscaled_size[0]}x{upscaled_size[1]}"
49
+
50
+ except Exception as e:
51
+
52
+ return None, None, "Error", str(e)
53
+
54
+
55
+ gr_interface = gr.Interface(
56
+ fn=gradio_interface,
57
+ inputs=gr.Image(type="pil"),
58
+ outputs=[gr.Image(), gr.Image(), gr.Text(), gr.Text()],
59
+ title="ESRGAN Image Upscaler",
60
+ description="Upload an image to upscale it using ESRGAN."
61
+ )
62
+
63
+ if __name__ == '__main__':
64
+
65
+ gr_interface.launch()