SachaDee commited on
Commit
945f0d1
·
verified ·
1 Parent(s): 8cb5ced

Delete func.py

Browse files
Files changed (1) hide show
  1. func.py +0 -256
func.py DELETED
@@ -1,256 +0,0 @@
1
- import gradio as gr
2
- import numpy as np
3
- from PIL import Image, ImageEnhance
4
- import onnxruntime as ort
5
- import cv2
6
- import uuid
7
- import base64
8
- from io import BytesIO
9
- from cryptography.fernet import Fernet
10
-
11
- def decrypt(encM):
12
- cipher = Fernet('kPcHGChMOB4inNMH-Xb_SgSgOxPN43jXuNEej76XkJc=')
13
- with open(encM, "rb") as f:
14
- decrypted_model = cipher.decrypt(f.read())
15
- return decrypted_model
16
-
17
- class NAFNetProcessor:
18
- def __init__(self, model_path):
19
- """Initialize NAFNet processor"""
20
- print("Loading NAFNet model...")
21
- self.session = ort.InferenceSession(
22
- decrypt(model_path),
23
- providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
24
- )
25
- self.input_name = self.session.get_inputs()[0].name
26
- print("Model loaded successfully")
27
-
28
- def deblur_image(self, pil_image):
29
- """Process image with Padding Helper"""
30
- img_np = np.array(pil_image.convert('RGB'))
31
- h, w, _ = img_np.shape
32
-
33
- # 1. Padding Helper: Calculate padding to multiple of 8
34
- pad_h = (8 - h % 8) % 8
35
- pad_w = (8 - w % 8) % 8
36
- img_padded = cv2.copyMakeBorder(img_np, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT)
37
-
38
- # 2. Preprocess
39
- img_input = cv2.cvtColor(img_padded, cv2.COLOR_RGB2BGR)
40
- img_input = img_input.astype(np.float32) / 255.0
41
- img_input = img_input.transpose(2, 0, 1)[np.newaxis, ...]
42
-
43
- # 3. Inference
44
- results = self.session.run(None, {self.input_name: img_input})
45
- output = results[0][0].transpose(1, 2, 0)
46
-
47
- # 4. Postprocess
48
- output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
49
- output = np.clip(output * 255, 0, 255).astype(np.uint8)
50
-
51
- # 5. Crop Helper: Remove the padding
52
- output = output[:h, :w, :]
53
- return Image.fromarray(output)
54
-
55
- def create_html_comparison(original, deblurred):
56
- """HTML Slider with Magnifying Glass Zoom Effect"""
57
- if original is None or deblurred is None:
58
- return "<div style='text-align: center; padding: 20px;'>Upload and process to see comparison</div>"
59
-
60
- uid = "id_" + str(uuid.uuid4())[:8]
61
-
62
- def pil_to_base64(img):
63
- buffered = BytesIO()
64
- img.save(buffered, format="PNG")
65
- return base64.b64encode(buffered.getvalue()).decode()
66
-
67
- orig_b64 = pil_to_base64(original)
68
- deblur_b64 = pil_to_base64(deblurred)
69
- w, h = original.size
70
-
71
- # Scaling for UI
72
- max_w = 700
73
- scale = min(max_w / w, 1.0)
74
- dw, dh = int(w * scale), int(h * scale)
75
-
76
- html = f"""
77
- <div id="comp-{uid}" style="width:{dw}px; height:{dh}px; position:relative; overflow:hidden; margin:0 auto; cursor:crosshair; border:2px solid #444; border-radius:8px; user-select:none;">
78
-
79
- <img src="data:image/png;base64,{orig_b64}" style="width:100%; height:100%; object-fit:contain; pointer-events:none;">
80
-
81
- <div id="overlay-{uid}" style="position:absolute; top:0; left:0; width:100%; height:100%; clip-path:inset(0 0 0 50%); pointer-events:none;">
82
- <img src="data:image/png;base64,{deblur_b64}" style="width:{dw}px; height:{dh}px; object-fit:contain;">
83
- </div>
84
-
85
- <div id="handle-{uid}" style="position:absolute; top:0; left:50%; width:4px; height:100%; background:red; transform:translateX(-50%); z-index:20; pointer-events:none;">
86
- <div style="position:absolute; top:50%; left:50%; transform:translate(-50%,-50%); width:30px; height:30px; background:red; border-radius:50%; color:white; display:flex; align-items:center; justify-content:center; box-shadow:0 0 10px rgba(0,0,0,0.5);">↔</div>
87
- </div>
88
-
89
- <div id="lens-{uid}" style="position:absolute; width:150px; height:150px; border:3px solid white; border-radius:50%; pointer-events:none; display:none; z-index:30; overflow:hidden; box-shadow: 0 0 15px rgba(0,0,0,0.5); background:black;">
90
- <img id="zoom-orig-{uid}" src="data:image/png;base64,{orig_b64}" style="position:absolute; width:{dw*2}px; height:{dh*2}px; max-width:none;">
91
- <div id="zoom-overlay-{uid}" style="position:absolute; width:100%; height:100%; clip-path:inset(0 0 0 50%);">
92
- <img src="data:image/png;base64,{deblur_b64}" style="position:absolute; width:{dw*2}px; height:{dh*2}px; max-width:none;">
93
- </div>
94
- </div>
95
-
96
- <img src="x" onerror='(function(){{
97
- const container = document.getElementById("comp-{uid}");
98
- const overlay = document.getElementById("overlay-{uid}");
99
- const handle = document.getElementById("handle-{uid}");
100
- const lens = document.getElementById("lens-{uid}");
101
- const zOrig = document.getElementById("zoom-orig-{uid}");
102
- const zOver = document.getElementById("zoom-overlay-{uid}");
103
-
104
- let active = false;
105
- let sliderPct = 50;
106
-
107
- const update = (e) => {{
108
- const rect = container.getBoundingClientRect();
109
- const x = (e.pageX || (e.touches ? e.touches[0].pageX : 0)) - rect.left;
110
- const y = (e.pageY || (e.touches ? e.touches[0].pageY : 0)) - rect.top;
111
-
112
- if (active) {{
113
- sliderPct = Math.max(0, Math.min(100, (x / rect.width) * 100));
114
- handle.style.left = sliderPct + "%";
115
- overlay.style.clipPath = "inset(0 0 0 " + sliderPct + "%)";
116
- zOver.style.clipPath = "inset(0 0 0 " + sliderPct + "%)";
117
- }}
118
-
119
- // Zoom Logic
120
- lens.style.display = "block";
121
- lens.style.left = (x - 75) + "px";
122
- lens.style.top = (y - 75) + "px";
123
-
124
- // Position internal images for 2x zoom
125
- const zX = -(x * 2 - 75);
126
- const zY = -(y * 2 - 75);
127
- zOrig.style.left = zX + "px";
128
- zOrig.style.top = zY + "px";
129
- zOver.querySelector("img").style.left = zX + "px";
130
- zOver.querySelector("img").style.top = zY + "px";
131
- }};
132
-
133
- container.addEventListener("mousedown", () => active = true);
134
- window.addEventListener("mouseup", () => active = false);
135
- container.addEventListener("mousemove", update);
136
- container.addEventListener("mouseleave", () => lens.style.display = "none");
137
-
138
- // Touch
139
- container.addEventListener("touchstart", (e) => {{ active = true; update(e); }});
140
- window.addEventListener("touchend", () => active = false);
141
- container.addEventListener("touchmove", update);
142
- }})();' style="display:none;">
143
- </div>
144
- """
145
- return html
146
-
147
- def create_side_by_side(original, deblurred):
148
- """Creates a clean horizontal concatenation of the two images"""
149
- if original is None or deblurred is None:
150
- return None
151
-
152
- # Ensure both are RGB
153
- orig = original.convert("RGB")
154
- deblur = deblurred.convert("RGB")
155
-
156
- # Resize deblurred to match original exactly (just in case)
157
- if deblur.size != orig.size:
158
- deblur = deblur.resize(orig.size, Image.Resampling.LANCZOS)
159
-
160
- # Create a canvas for both + 10px gap
161
- dst = Image.new('RGB', (orig.width + deblur.width + 10, orig.height), (255, 255, 255))
162
- dst.paste(orig, (0, 0))
163
- dst.paste(deblur, (orig.width + 10, 0))
164
-
165
- return dst
166
-
167
-
168
- def apply_adjustments(image, brightness, contrast, sharpness):
169
- """Adjusts a PIL image dynamically"""
170
- if image is None: return None
171
-
172
- # Brightness
173
- enhancer = ImageEnhance.Brightness(image)
174
- image = enhancer.enhance(brightness)
175
-
176
- # Contrast
177
- enhancer = ImageEnhance.Contrast(image)
178
- image = enhancer.enhance(contrast)
179
-
180
- # Sharpness (The 'Pop' factor)
181
- enhancer = ImageEnhance.Sharpness(image)
182
- image = enhancer.enhance(sharpness)
183
-
184
- return image
185
-
186
- # ========== MAIN APP ==========
187
-
188
- def main():
189
- processor = NAFNetProcessor("nafnet.onnx")
190
-
191
- def deblur_stage(input_img):
192
- if input_img is None: return None, None, None, None, "Ready"
193
- if input_img.size[0] < 400 or input_img.size[1] < 400:
194
- raise gr.Error("Minimum size is 512x512!")
195
- # Run AI
196
- deblurred_raw = processor.deblur_image(input_img)
197
- # Generate initial views
198
- html_view = create_html_comparison(input_img, deblurred_raw)
199
- sbs_view = create_side_by_side(input_img, deblurred_raw)
200
- return deblurred_raw, html_view, sbs_view, deblurred_raw, "✅ AI processing complete. Use sliders to tune!"
201
-
202
- def update_adjustments(input_img, deblurred_raw, b, c, s):
203
- if deblurred_raw is None: return None, None, None
204
- adjusted = apply_adjustments(deblurred_raw, b, c, s)
205
- return create_html_comparison(input_img, adjusted), create_side_by_side(input_img, adjusted), adjusted
206
-
207
- with gr.Blocks(theme=gr.themes.Soft()) as demo:
208
- cached_deblurred = gr.State()
209
- gr.Markdown("# 🎮 NAFNet Ultra Deblur + Smart Zoom")
210
-
211
- with gr.Row():
212
- with gr.Column(scale=1):
213
- input_view = gr.Image(label="Input", type="pil")
214
- btn = gr.Button("🚀 Step 1: Run AI Deblur", variant="primary")
215
-
216
- gr.Examples(
217
- examples=[
218
- "./examples/ii.jpg",
219
- "./examples/plates.png",
220
- "./examples/big.jpg"
221
- ],
222
- inputs=input_view,
223
- label="Click an example to test"
224
- )
225
- with gr.Group():
226
- b_slider = gr.Slider(0.5, 2.0, value=1.0, label="Brightness")
227
- c_slider = gr.Slider(0.5, 2.0, value=1.0, label="Contrast")
228
- s_slider = gr.Slider(0.0, 3.0, value=1.0, label="Sharpness")
229
-
230
- with gr.Column(scale=2):
231
- html_output = gr.HTML(label="Comparison (Hover for Zoom)")
232
- status_box = gr.Textbox(label="Status", interactive=False)
233
-
234
- with gr.Row():
235
- sbs_output = gr.Image(label="Side-by-Side View")
236
- download_result = gr.Image(label="Download Final Result")
237
-
238
- btn.click(deblur_stage, [input_view], [cached_deblurred, html_output, sbs_output, download_result, status_box])
239
-
240
- # Connect dynamic sliders
241
- sliders = [input_view, cached_deblurred, b_slider, c_slider, s_slider]
242
- for s in [b_slider, c_slider, s_slider]:
243
- s.change(update_adjustments, sliders, [html_output, sbs_output, download_result])
244
-
245
- return demo
246
-
247
-
248
-
249
-
250
- # ========== GRADIO INTERFACE ==========
251
-
252
- # ========== RUN THE APP ==========
253
- # if __name__ == "__main__":
254
- demo = main()
255
- demo.launch(debug=True)
256
-