Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image, ImageEnhance | |
| import onnxruntime as ort | |
| import cv2 | |
| import uuid | |
| import base64 | |
| from io import BytesIO | |
| from cryptography.fernet import Fernet | |
| def decrypt(encM): | |
| cipher = Fernet('kPcHGChMOB4inNMH-Xb_SgSgOxPN43jXuNEej76XkJc=') | |
| with open(encM, "rb") as f: | |
| decrypted_model = cipher.decrypt(f.read()) | |
| return decrypted_model | |
| class NAFNetProcessor: | |
| def __init__(self, model_path): | |
| """Initialize NAFNet processor""" | |
| print("Loading NAFNet model...") | |
| self.session = ort.InferenceSession( | |
| decrypt(model_path), | |
| providers=['CUDAExecutionProvider', 'CPUExecutionProvider'] | |
| ) | |
| self.input_name = self.session.get_inputs()[0].name | |
| print("Model loaded successfully") | |
| def deblur_image(self, pil_image): | |
| """Process image with Padding Helper""" | |
| img_np = np.array(pil_image.convert('RGB')) | |
| h, w, _ = img_np.shape | |
| # 1. Padding Helper: Calculate padding to multiple of 8 | |
| pad_h = (8 - h % 8) % 8 | |
| pad_w = (8 - w % 8) % 8 | |
| img_padded = cv2.copyMakeBorder(img_np, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT) | |
| # 2. Preprocess | |
| img_input = cv2.cvtColor(img_padded, cv2.COLOR_RGB2BGR) | |
| img_input = img_input.astype(np.float32) / 255.0 | |
| img_input = img_input.transpose(2, 0, 1)[np.newaxis, ...] | |
| # 3. Inference | |
| results = self.session.run(None, {self.input_name: img_input}) | |
| output = results[0][0].transpose(1, 2, 0) | |
| # 4. Postprocess | |
| output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB) | |
| output = np.clip(output * 255, 0, 255).astype(np.uint8) | |
| # 5. Crop Helper: Remove the padding | |
| output = output[:h, :w, :] | |
| return Image.fromarray(output) | |
| def create_html_comparison(original, deblurred): | |
| """HTML Slider with Magnifying Glass Zoom Effect""" | |
| if original is None or deblurred is None: | |
| return "<div style='text-align: center; padding: 20px;'>Upload and process to see comparison</div>" | |
| uid = "id_" + str(uuid.uuid4())[:8] | |
| def pil_to_base64(img): | |
| buffered = BytesIO() | |
| img.save(buffered, format="PNG") | |
| return base64.b64encode(buffered.getvalue()).decode() | |
| orig_b64 = pil_to_base64(original) | |
| deblur_b64 = pil_to_base64(deblurred) | |
| w, h = original.size | |
| # Scaling for UI | |
| max_w = 700 | |
| scale = min(max_w / w, 1.0) | |
| dw, dh = int(w * scale), int(h * scale) | |
| html = f""" | |
| <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;"> | |
| <img src="data:image/png;base64,{orig_b64}" style="width:100%; height:100%; object-fit:contain; pointer-events:none;"> | |
| <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;"> | |
| <img src="data:image/png;base64,{deblur_b64}" style="width:{dw}px; height:{dh}px; object-fit:contain;"> | |
| </div> | |
| <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;"> | |
| <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> | |
| </div> | |
| <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;"> | |
| <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;"> | |
| <div id="zoom-overlay-{uid}" style="position:absolute; width:100%; height:100%; clip-path:inset(0 0 0 50%);"> | |
| <img src="data:image/png;base64,{deblur_b64}" style="position:absolute; width:{dw*2}px; height:{dh*2}px; max-width:none;"> | |
| </div> | |
| </div> | |
| <img src="x" onerror='(function(){{ | |
| const container = document.getElementById("comp-{uid}"); | |
| const overlay = document.getElementById("overlay-{uid}"); | |
| const handle = document.getElementById("handle-{uid}"); | |
| const lens = document.getElementById("lens-{uid}"); | |
| const zOrig = document.getElementById("zoom-orig-{uid}"); | |
| const zOver = document.getElementById("zoom-overlay-{uid}"); | |
| let active = false; | |
| let sliderPct = 50; | |
| const update = (e) => {{ | |
| const rect = container.getBoundingClientRect(); | |
| const x = (e.pageX || (e.touches ? e.touches[0].pageX : 0)) - rect.left; | |
| const y = (e.pageY || (e.touches ? e.touches[0].pageY : 0)) - rect.top; | |
| if (active) {{ | |
| sliderPct = Math.max(0, Math.min(100, (x / rect.width) * 100)); | |
| handle.style.left = sliderPct + "%"; | |
| overlay.style.clipPath = "inset(0 0 0 " + sliderPct + "%)"; | |
| zOver.style.clipPath = "inset(0 0 0 " + sliderPct + "%)"; | |
| }} | |
| // Zoom Logic | |
| lens.style.display = "block"; | |
| lens.style.left = (x - 75) + "px"; | |
| lens.style.top = (y - 75) + "px"; | |
| // Position internal images for 2x zoom | |
| const zX = -(x * 2 - 75); | |
| const zY = -(y * 2 - 75); | |
| zOrig.style.left = zX + "px"; | |
| zOrig.style.top = zY + "px"; | |
| zOver.querySelector("img").style.left = zX + "px"; | |
| zOver.querySelector("img").style.top = zY + "px"; | |
| }}; | |
| container.addEventListener("mousedown", () => active = true); | |
| window.addEventListener("mouseup", () => active = false); | |
| container.addEventListener("mousemove", update); | |
| container.addEventListener("mouseleave", () => lens.style.display = "none"); | |
| // Touch | |
| container.addEventListener("touchstart", (e) => {{ active = true; update(e); }}); | |
| window.addEventListener("touchend", () => active = false); | |
| container.addEventListener("touchmove", update); | |
| }})();' style="display:none;"> | |
| </div> | |
| """ | |
| return html | |
| def create_side_by_side(original, deblurred): | |
| """Creates a clean horizontal concatenation of the two images""" | |
| if original is None or deblurred is None: | |
| return None | |
| # Ensure both are RGB | |
| orig = original.convert("RGB") | |
| deblur = deblurred.convert("RGB") | |
| # Resize deblurred to match original exactly (just in case) | |
| if deblur.size != orig.size: | |
| deblur = deblur.resize(orig.size, Image.Resampling.LANCZOS) | |
| # Create a canvas for both + 10px gap | |
| dst = Image.new('RGB', (orig.width + deblur.width + 10, orig.height), (255, 255, 255)) | |
| dst.paste(orig, (0, 0)) | |
| dst.paste(deblur, (orig.width + 10, 0)) | |
| return dst | |
| def apply_adjustments(image, brightness, contrast, sharpness): | |
| """Adjusts a PIL image dynamically""" | |
| if image is None: return None | |
| # Brightness | |
| enhancer = ImageEnhance.Brightness(image) | |
| image = enhancer.enhance(brightness) | |
| # Contrast | |
| enhancer = ImageEnhance.Contrast(image) | |
| image = enhancer.enhance(contrast) | |
| # Sharpness (The 'Pop' factor) | |
| enhancer = ImageEnhance.Sharpness(image) | |
| image = enhancer.enhance(sharpness) | |
| return image | |
| # ========== MAIN APP ========== | |
| def main(): | |
| processor = NAFNetProcessor("nafnet.onnx") | |
| def deblur_stage(input_img): | |
| if input_img is None: return None, None, None, None, "Ready" | |
| if input_img.size[0] < 400 or input_img.size[1] < 400: | |
| raise gr.Error("Minimum size is 512x512!") | |
| # Run AI | |
| deblurred_raw = processor.deblur_image(input_img) | |
| # Generate initial views | |
| html_view = create_html_comparison(input_img, deblurred_raw) | |
| sbs_view = create_side_by_side(input_img, deblurred_raw) | |
| return deblurred_raw, html_view, sbs_view, deblurred_raw, "✅ AI processing complete. Use sliders to tune!" | |
| def update_adjustments(input_img, deblurred_raw, b, c, s): | |
| if deblurred_raw is None: return None, None, None | |
| adjusted = apply_adjustments(deblurred_raw, b, c, s) | |
| return create_html_comparison(input_img, adjusted), create_side_by_side(input_img, adjusted), adjusted | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| cached_deblurred = gr.State() | |
| gr.Markdown("# 🎮 NAFNet Ultra Deblur + Smart Zoom") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| input_view = gr.Image(label="Input", type="pil") | |
| btn = gr.Button("🚀 Step 1: Run AI Deblur", variant="primary") | |
| gr.Examples( | |
| examples=[ | |
| "./examples/ii.jpg", | |
| "./examples/plates.png", | |
| "./examples/big.jpg" | |
| ], | |
| inputs=input_view, | |
| label="Click an example to test" | |
| ) | |
| with gr.Group(): | |
| b_slider = gr.Slider(0.5, 2.0, value=1.0, label="Brightness") | |
| c_slider = gr.Slider(0.5, 2.0, value=1.0, label="Contrast") | |
| s_slider = gr.Slider(0.0, 3.0, value=1.0, label="Sharpness") | |
| with gr.Column(scale=2): | |
| html_output = gr.HTML(label="Comparison (Hover for Zoom)") | |
| status_box = gr.Textbox(label="Status", interactive=False) | |
| with gr.Row(): | |
| sbs_output = gr.Image(label="Side-by-Side View") | |
| download_result = gr.Image(label="Download Final Result") | |
| btn.click(deblur_stage, [input_view], [cached_deblurred, html_output, sbs_output, download_result, status_box]) | |
| # Connect dynamic sliders | |
| sliders = [input_view, cached_deblurred, b_slider, c_slider, s_slider] | |
| for s in [b_slider, c_slider, s_slider]: | |
| s.change(update_adjustments, sliders, [html_output, sbs_output, download_result]) | |
| return demo | |
| # ========== GRADIO INTERFACE ========== | |
| with gr.Blocks() as demo: | |
| gr.Markdown(""" | |
| # 🎮 NAFNet Image Deblurring | |
| **Upload a blurry image and compare with deblurred result** | |
| """) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| # Upload | |
| upload = gr.Image( | |
| label="📤 Upload Image", | |
| type="pil", | |
| height=250 | |
| ) | |
| # Process button | |
| process_btn = gr.Button( | |
| "🚀 Process Image", | |
| variant="primary", | |
| size="lg" | |
| ) | |
| # Status | |
| status = gr.Textbox( | |
| label="Status", | |
| value="Ready to process", | |
| interactive=False | |
| ) | |
| gr.Markdown(""" | |
| ### Instructions: | |
| 1. Upload blurry image | |
| 2. Click **Process Image** | |
| 3. Wait a few seconds | |
| 4. Drag the red line to compare | |
| """) | |
| with gr.Column(scale=2): | |
| # Interactive comparison | |
| html_output = gr.HTML( | |
| label="🔄 Interactive Comparison", | |
| value="<div style='text-align: center; padding: 50px; color: #666;'>Upload an image to begin</div>" | |
| ) | |
| # Side-by-side | |
| with gr.Accordion("📸 Side-by-Side View", open=True): | |
| side_by_side = gr.Image( | |
| label="Original vs Deblurred", | |
| type="pil", | |
| height=350 | |
| ) | |
| # Event handlers | |
| process_btn.click( | |
| fn=process_image, | |
| inputs=[upload], | |
| outputs=[html_output, side_by_side, status] | |
| ) | |
| # Clear when new image uploaded | |
| upload.change( | |
| fn=lambda: ( | |
| "<div style='text-align: center; padding: 50px; color: #666;'>Click 'Process Image' to start</div>", | |
| None, | |
| "Image ready for processing" | |
| ), | |
| inputs=[], | |
| outputs=[html_output, side_by_side, status] | |
| ) | |
| return demo | |
| # ========== RUN THE APP ========== | |
| # if __name__ == "__main__": | |
| demo = main() | |
| demo.launch(debug=True) | |