Update app.py

app.py

@@ -1,113 +1,67 @@
-import gradio as gr
 import torch
 import torch.nn as nn
-from transformers import CLIPProcessor, CLIPModel
+import torchvision.transforms as transforms
 from PIL import Image
 import numpy as np
+import gradio as gr
 
-# -----------------------------
-# 1. Zero-DCE model (light enhancement)
-# -----------------------------
+# -----------------------
+# Zero-DCE Model
+# -----------------------
 class ZeroDCE(nn.Module):
-    def __init__(self):
+    def __init__(self, num_layers=7):
         super(ZeroDCE, self).__init__()
-        self.relu = nn.ReLU(inplace=True)
-        self.conv1 = nn.Conv2d(3, 32, 3, padding=1, bias=True)
-        self.conv2 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
-        self.conv3 = nn.Conv2d(32, 32, 3, padding=1, bias=True)
-        self.conv4 = nn.Conv2d(32, 24, 3, padding=1, bias=True)
-        self.conv5 = nn.Conv2d(24, 8, 3, padding=1, bias=True)
+        filters = 32
+        layers = []
+        layers += [nn.Conv2d(3, filters, 3, 1, 1), nn.ReLU(inplace=True)]
+        for _ in range(num_layers - 2):
+            layers += [nn.Conv2d(filters, filters, 3, 1, 1), nn.ReLU(inplace=True)]
+        layers += [nn.Conv2d(filters, 3, 3, 1, 1), nn.Tanh()]
+        self.net = nn.Sequential(*layers)
 
     def forward(self, x):
-        x1 = self.relu(self.conv1(x))
-        x2 = self.relu(self.conv2(x1))
-        x3 = self.relu(self.conv3(x2))
-        x4 = self.relu(self.conv4(x3))
-        out = torch.tanh(self.conv5(x4))
-        return out
-
-def enhance_image(img, model):
-    # Convert PIL -> Tensor
-    img_np = np.array(img).astype(np.float32) / 255.0
-    img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).to(device)
-
-    with torch.no_grad():
-        enhancement_map = model(img_tensor)
+        return x + self.net(x)
 
-    # Apply enhancement: simple iterative curve
-    enhanced = img_tensor
-    for i in range(8):
-        enhanced = enhanced + enhancement_map[:, i*3:(i+1)*3, :, :] * (enhanced**2 - enhanced)
-
-    enhanced = torch.clamp(enhanced, 0, 1)
-    enhanced = enhanced.squeeze(0).permute(1, 2, 0).cpu().numpy()
-    enhanced = (enhanced * 255).astype(np.uint8)
-
-    return Image.fromarray(enhanced)
+# -----------------------
+# Setup
+# -----------------------
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+zero_dce = ZeroDCE().to(device)
 
-# -----------------------------
-# 2. Aesthetic Scoring Model
-# -----------------------------
-class AestheticPredictor(nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.clip = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
-        self.mlp = nn.Sequential(
-            nn.Linear(self.clip.config.projection_dim, 512),
-            nn.ReLU(),
-            nn.Linear(512, 1)
-        )
+# -----------------------
+# Image Enhancement
+# -----------------------
+def enhance_image(image, model):
+    # Always force RGB
+    image = image.convert("RGB")
 
-    def forward(self, pixel_values, input_ids, attention_mask):
-        outputs = self.clip(pixel_values=pixel_values, input_ids=input_ids, attention_mask=attention_mask)
-        pooled_output = outputs.pooler_output
-        return self.mlp(pooled_output)
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+    ])
+    img_tensor = transform(image).unsqueeze(0).to(device)
 
-def score_image(image, processor, model):
-    inputs = processor(text=["aesthetic photo"], images=image, return_tensors="pt", padding=True).to(device)
     with torch.no_grad():
-        score = model(**inputs)
-    return score.item()
+        enhanced = model(img_tensor)
 
-# -----------------------------
-# 3. Pipeline function
-# -----------------------------
-def process_image(input_img):
-    # Step 1: enhance
-    enhanced_img = enhance_image(input_img, zero_dce)
-
-    # Step 2: aesthetic scoring
-    original_score = score_image(input_img, processor, ae_model)
-    enhanced_score = score_image(enhanced_img, processor, ae_model)
-
-    # Step 3: choose best
-    if enhanced_score > original_score:
-        return enhanced_img, f"Enhanced chosen (score {enhanced_score:.2f} vs {original_score:.2f})"
-    else:
-        return input_img, f"Original kept (score {original_score:.2f} vs {enhanced_score:.2f})"
+    # Convert back to PIL
+    enhanced = enhanced.squeeze(0).cpu().permute(1, 2, 0).numpy()
+    enhanced = (enhanced * 255).clip(0, 255).astype(np.uint8)
 
-# -----------------------------
-# 4. Setup
-# -----------------------------
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return Image.fromarray(enhanced)
 
-zero_dce = ZeroDCE().to(device)
-ae_model = AestheticPredictor().to(device)
-processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")
+# -----------------------
+# Gradio Interface
+# -----------------------
+def process_image(input_img):
+    return enhance_image(input_img, zero_dce)
 
-# -----------------------------
-# 5. Gradio UI
-# -----------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## 📸 AI Photo Enhancer with Aesthetic Scoring")
-
+    gr.Markdown("## 🌙 Low Light Image Enhancement (Zero-DCE)")
     with gr.Row():
-        inp = gr.Image(type="pil", label="Upload your photo")
-        out = gr.Image(type="pil", label="Best looking result")
-    
-    info = gr.Label(label="Result Info")
-
-    btn = gr.Button("Enhance ✨")
-    btn.click(process_image, inputs=inp, outputs=[out, info])
+        inp = gr.Image(type="pil", label="Upload Dark Image")
+        out = gr.Image(type="pil", label="Enhanced Image")
+    run_btn = gr.Button("Enhance")
+    run_btn.click(fn=process_image, inputs=inp, outputs=out)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)