1w33

.gitignore

@@ -0,0 +1,54 @@
+# macOS system files
+.DS_Store
+
+# Node.js dependencies
+node_modules/
+
+# Python
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+env/
+venv/
+ENV/
+.venv/
+
+# Logs
+*.log
+
+# OS generated files
+Thumbs.db
+ehthumbs.db
+Icon?
+Desktop.ini
+
+# IDEs and editors
+.vscode/
+.idea/
+*.sublime-workspace
+*.sublime-project
+
+# Build output
+dist/
+build/
+out/
+
+# Temporary files
+*.tmp
+*.swp
+*.bak
+*.orig
+
+# Test coverage
+coverage/
+*.cover
+.nyc_output/
+
+# Environment files
+.env
+.env.*
+
+# Yarn
+.yarn/
+.pnp.*

app.py

@@ -0,0 +1,299 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import gradio as gr
+import numpy as np
+from PIL import Image
+import torchvision.transforms as tfs
+import os
+
+def default_conv(in_channels, out_channels, kernel_size, bias=True):
+    return nn.Conv2d(in_channels, out_channels, kernel_size, padding=(kernel_size // 2), bias=bias)
+
+class PALayer(nn.Module):
+    def __init__(self, channel):
+        super(PALayer, self).__init__()
+        self.pa = nn.Sequential(
+            nn.Conv2d(channel, channel // 8, 1, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channel // 8, 1, 1, bias=True),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        y = self.pa(x)
+        return x * y
+
+class CALayer(nn.Module):
+    def __init__(self, channel):
+        super(CALayer, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.ca = nn.Sequential(
+            nn.Conv2d(channel, channel // 8, 1, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(channel // 8, channel, 1, bias=True),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        y = self.avg_pool(x)
+        y = self.ca(y)
+        return x * y
+
+class Block(nn.Module):
+    def __init__(self, conv, dim, kernel_size):
+        super(Block, self).__init__()
+        self.conv1 = conv(dim, dim, kernel_size, bias=True)
+        self.act1 = nn.ReLU(inplace=True)
+        self.conv2 = conv(dim, dim, kernel_size, bias=True)
+        self.calayer = CALayer(dim)
+        self.palayer = PALayer(dim)
+
+    def forward(self, x):
+        res = self.act1(self.conv1(x))
+        res = res + x
+        res = self.conv2(res)
+        res = self.calayer(res)
+        res = self.palayer(res)
+        res += x
+        return res
+
+class Group(nn.Module):
+    def __init__(self, conv, dim, kernel_size, blocks):
+        super(Group, self).__init__()
+        modules = [Block(conv, dim, kernel_size) for _ in range(blocks)]
+        modules.append(conv(dim, dim, kernel_size))
+        self.gp = nn.Sequential(*modules)
+
+    def forward(self, x):
+        res = self.gp(x)
+        res += x
+        return res
+
+class FFA(nn.Module):
+    def __init__(self, gps, blocks, conv=default_conv):
+        super(FFA, self).__init__()
+        self.gps = gps
+        self.dim = 64
+        kernel_size = 3
+
+        pre_process = [conv(3, self.dim, kernel_size)]
+        assert self.gps == 3
+        self.g1 = Group(conv, self.dim, kernel_size, blocks=blocks)
+        self.g2 = Group(conv, self.dim, kernel_size, blocks=blocks)
+        self.g3 = Group(conv, self.dim, kernel_size, blocks=blocks)
+        self.ca = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            nn.Conv2d(self.dim * self.gps, self.dim // 16, 1, bias=True),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(self.dim // 16, self.dim * self.gps, 1, bias=True),
+            nn.Sigmoid()
+        )
+        self.palayer = PALayer(self.dim)
+
+        post_process = [
+            conv(self.dim, self.dim, kernel_size),
+            conv(self.dim, 3, kernel_size)
+        ]
+
+        self.pre = nn.Sequential(*pre_process)
+        self.post = nn.Sequential(*post_process)
+
+    def forward(self, x1):
+        x = self.pre(x1)
+        res1 = self.g1(x)
+        res2 = self.g2(res1)
+        res3 = self.g3(res2)
+        w = self.ca(torch.cat([res1, res2, res3], dim=1))
+        w = w.view(-1, self.gps, self.dim)[:, :, :, None, None]
+        out = w[:, 0, :, :, :] * res1 + w[:, 1, :, :, :] * res2 + w[:, 2, :, :, :] * res3
+        out = self.palayer(out)
+        x = self.post(out)
+        return x + x1
+
+MODEL_PATH = 'tti.pk'
+gps = 3
+blocks = 19
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+net = FFA(gps=gps, blocks=blocks).to(device)
+net = torch.nn.DataParallel(net)
+
+if not os.path.exists(MODEL_PATH):
+    raise FileNotFoundError(f"Model checkpoint not found at {MODEL_PATH}")
+
+try:
+
+    torch.serialization.add_safe_globals([np.core.multiarray.scalar])
+    checkpoint = torch.load(MODEL_PATH, map_location=device, weights_only=True)
+except:
+
+    print("Warning: Loading checkpoint with weights_only=False. Ensure the checkpoint is from a trusted source.")
+    checkpoint = torch.load(MODEL_PATH, map_location=device, weights_only=False)
+net.load_state_dict(checkpoint['model'])
+net.eval()
+
+print(f"Model loaded successfully on {device}")
+
+def dehaze_image(image):
+    """
+    Process a hazy image and return the dehazed result.
+
+    Args:
+        image: PIL Image or numpy array
+
+    Returns:
+        PIL Image: Dehazed image
+    """
+    try:
+
+        if isinstance(image, np.ndarray):
+            image = Image.fromarray(image)
+
+
+        haze_img = image.convert("RGB")
+
+
+        transform = tfs.Compose([
+            tfs.ToTensor(),
+            tfs.Normalize(mean=[0.64, 0.6, 0.58], std=[0.14, 0.15, 0.152])
+        ])
+
+        haze_tensor = transform(haze_img).unsqueeze(0).to(device)
+
+
+        with torch.no_grad():
+            pred = net(haze_tensor)
+
+
+        pred_clamped = pred.clamp(0, 1).cpu()
+        pred_numpy = pred_clamped.squeeze(0).permute(1, 2, 0).numpy()
+        pred_numpy = (pred_numpy * 255).astype(np.uint8)
+
+        return Image.fromarray(pred_numpy)
+
+    except Exception as e:
+        print(f"Error processing image: {str(e)}")
+        return None
+
+SAMPLE_IMAGES = [
+    "./img/s2.png",
+    "./img/s4.png"
+]
+
+def load_sample_image(sample_path):
+    """Load and return a sample image"""
+    try:
+        if os.path.exists(sample_path):
+            return Image.open(sample_path)
+        else:
+            print(f"Sample image not found: {sample_path}")
+            return None
+    except Exception as e:
+        print(f"Error loading sample image {sample_path}: {e}")
+        return None
+
+def create_interface():
+    with gr.Blocks(title="Image Dehazing App", theme=gr.themes.Soft()) as demo:
+        gr.Markdown("# 🌫️ Image Dehazing with FFA-Net")
+        gr.Markdown("Upload a hazy image to remove fog, haze, and improve visibility!")
+
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(
+                    label="Upload Hazy Image",
+                    type="pil",
+                    height=400
+                )
+
+
+                gr.Markdown("### Try Sample Images")
+                with gr.Row():
+                    sample1_btn = gr.Image(
+                        value=load_sample_image(SAMPLE_IMAGES[0]) if len(SAMPLE_IMAGES) > 0 else None,
+                        label="Sample 1",
+                        interactive=True,
+                        width=150,
+                        height=150,
+                        container=True,
+                        show_download_button=False
+                    )
+                    sample2_btn = gr.Image(
+                        value=load_sample_image(SAMPLE_IMAGES[1]) if len(SAMPLE_IMAGES) > 1 else None,
+                        label="Sample 2",
+                        interactive=True,
+                        width=150,
+                        height=150,
+                        container=True,
+                        show_download_button=False
+                    )
+
+                process_btn = gr.Button(
+                    "Remove Haze ✨",
+                    variant="primary",
+                    size="lg"
+                )
+
+            with gr.Column():
+                output_image = gr.Image(
+                    label="Dehazed Result",
+                    type="pil",
+                    height=400
+                )
+
+
+        def use_sample1():
+            return load_sample_image(SAMPLE_IMAGES[0]) if len(SAMPLE_IMAGES) > 0 else None
+
+        def use_sample2():
+            return load_sample_image(SAMPLE_IMAGES[1]) if len(SAMPLE_IMAGES) > 1 else None
+
+        sample1_btn.select(
+            fn=use_sample1,
+            outputs=input_image
+        )
+
+        sample2_btn.select(
+            fn=use_sample2,
+            outputs=input_image
+        )
+
+
+        process_btn.click(
+            fn=dehaze_image,
+            inputs=input_image,
+            outputs=output_image,
+            api_name="dehaze"
+        )
+
+
+        input_image.change(
+            fn=dehaze_image,
+            inputs=input_image,
+            outputs=output_image
+        )
+
+        gr.Markdown("""
+        ### About
+        This app uses the FFA-Net (Feature Fusion Attention Network) for single image dehazing.
+        The model removes atmospheric haze and fog to restore clear, vibrant images.
+
+        **Tips for best results:**
+        - Use good quality images with visible haze or fog
+        - Model works best on indoor images
+        **Made by  <a href="https://www.linkedin.com/in/aditsg26/">Aditya Singh</a> and <a href="https://www.linkedin.com/in/ramandeep-singh-makkar/">Ramandeep Singh Makkar</a>**
+        """)
+
+    return demo
+
+if __name__ == "__main__":
+
+    demo = create_interface()
+
+
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        debug=False
+    )

requirements.txt

@@ -0,0 +1,4 @@
+gradio
+pillow
+torch
+torchvision