Upload func.py

func.py

@@ -0,0 +1,327 @@
+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)
+