wtf741 commited on
Commit
2cd2c3f
1 Parent(s): bc62a9a

init commit

Browse files
Files changed (3) hide show
  1. .gitignore +4 -0
  2. main.py +68 -0
  3. requirements.txt +6 -0
.gitignore ADDED
@@ -0,0 +1,4 @@
 
 
 
 
 
1
+ /venv
2
+ __pycache__
3
+ /test_*
4
+ /.idea
main.py ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+
3
+ import cv2
4
+ import gradio as gr
5
+ import numpy as np
6
+ from PIL import Image
7
+ from cv2.ximgproc import guidedFilter
8
+ from imgutils.data import load_image
9
+ from imgutils.restore import restore_with_nafnet, restore_with_scunet
10
+
11
+
12
+ def clean_adverse(
13
+ input_image: Image.Image,
14
+ diameter: float = 5,
15
+ sigma_color: float = 8,
16
+ sigma_space: float = 8,
17
+ radius: float = 4,
18
+ eps: float = 16,
19
+ ) -> Image.Image:
20
+ img = np.array(input_image).astype(np.float32)
21
+ y = img.copy()
22
+
23
+ for _ in range(64):
24
+ y = cv2.bilateralFilter(y, diameter, sigma_color, sigma_space)
25
+
26
+ for _ in range(4):
27
+ y = guidedFilter(img, y, radius, eps)
28
+
29
+ output_image = Image.fromarray(y.clip(0, 255).astype(np.uint8))
30
+ return output_image
31
+
32
+
33
+ def clean(
34
+ image: Image.Image,
35
+ use_adverse_clean: bool = True,
36
+ use_scunet_clean: bool = True,
37
+ use_nafnet_clean: bool = False
38
+ ) -> Image.Image:
39
+ image = load_image(image)
40
+ if use_adverse_clean:
41
+ image = clean_adverse(image)
42
+ if use_scunet_clean:
43
+ image = restore_with_scunet(image)
44
+ if use_nafnet_clean:
45
+ image = restore_with_nafnet(image)
46
+ return image
47
+
48
+
49
+ if __name__ == '__main__':
50
+ with gr.Blocks() as demo:
51
+ with gr.Row():
52
+ with gr.Column():
53
+ gr_input_image = gr.Image(label='Input Image', type="pil")
54
+ gr_submit = gr.Button(value='FUCK YOU MIST!')
55
+ gr_adverse = gr.Checkbox(label='Use Adverse Clean', value=True)
56
+ gr_scunet = gr.Checkbox(label='Use SCUNET', value=True)
57
+ gr_nafnet = gr.Checkbox(label='Use NafNET', value=False)
58
+
59
+ with gr.Column():
60
+ gr_output_image = gr.Image(label='Output Image', type="pil")
61
+
62
+ gr_submit.click(
63
+ fn=clean,
64
+ inputs=[gr_input_image, gr_adverse, gr_scunet, gr_nafnet],
65
+ outputs=[gr_output_image],
66
+ )
67
+
68
+ demo.queue(os.cpu_count()).launch()
requirements.txt ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ gradio==4.10.0
2
+ pillow
3
+ dghs-imgutils
4
+ numpy
5
+ opencv-python
6
+ opencv-contrib-python