Spaces:

syedaoon
/

ZeroIG

Running

App Files Files Community

syedaoon commited on Jul 27

Commit

d59f3c4

verified ·

1 Parent(s): 5de8ab6

Update app.py

Browse files

Files changed (1) hide show

app.py +159 -33

app.py CHANGED Viewed

@@ -1,8 +1,11 @@
-from flask import Flask, request, render_template_string, send_file
 from PIL import Image
 import numpy as np
 import io
 import base64
 app = Flask(__name__)
@@ -12,76 +15,199 @@ HTML_TEMPLATE = """
 <head>
     <title>ZeroIG Enhancement</title>
     <style>
-        body { font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; padding: 20px; }
         .container { text-align: center; }
-        .upload-area { border: 2px dashed #ccc; padding: 40px; margin: 20px 0; }
         .result { margin-top: 20px; }
-        img { max-width: 100%; height: auto; }
     </style>
 </head>
 <body>
     <div class="container">
-        <h1>🌟 ZeroIG: Low-Light Enhancement</h1>
-        <p>Upload a low-light image to enhance it!</p>
         <form method="post" enctype="multipart/form-data">
             <div class="upload-area">
                 <input type="file" name="image" accept="image/*" required>
                 <br><br>
-                <button type="submit">Enhance Image</button>
             </div>
         </form>
-        {% if result_image %}
         <div class="result">
-            <h3>Enhanced Image:</h3>
-            <img src="data:image/png;base64,{{ result_image }}" alt="Enhanced">
-            <br><br>
-            <a href="data:image/png;base64,{{ result_image }}" download="enhanced.png">Download Enhanced Image</a>
         </div>
         {% endif %}
     </div>
 </body>
 </html>
 """
-def enhance_image(image):
-    """Simple image enhancement"""
-    try:
-        # Convert to numpy array
-        arr = np.array(image)
-        # Simple brightness enhancement
-        enhanced = np.clip(arr.astype(np.float32) * 1.8, 0, 255).astype(np.uint8)
         return Image.fromarray(enhanced)
-    except Exception as e:
-        print(f"Enhancement error: {e}")
-        return image
 @app.route('/', methods=['GET', 'POST'])
 def index():
     result_image = None
     if request.method == 'POST':
         try:
             file = request.files['image']
             if file:
                 # Open and process image
                 image = Image.open(file.stream).convert('RGB')
-                enhanced = enhance_image(image)
-                # Convert to base64 for display
-                img_buffer = io.BytesIO()
-                enhanced.save(img_buffer, format='PNG')
-                img_str = base64.b64encode(img_buffer.getvalue()).decode()
-                result_image = img_str
         except Exception as e:
-            print(f"Error processing image: {e}")
-    return render_template_string(HTML_TEMPLATE, result_image=result_image)
 if __name__ == '__main__':
-    print("🚀 Starting Flask app...")
     app.run(host='0.0.0.0', port=7860)

+from flask import Flask, request, render_template_string
 from PIL import Image
 import numpy as np
 import io
 import base64
+import torch
+import torchvision.transforms as transforms
+import os
 app = Flask(__name__)
 <head>
     <title>ZeroIG Enhancement</title>
     <style>
+        body { font-family: Arial, sans-serif; max-width: 1000px; margin: 0 auto; padding: 20px; }
         .container { text-align: center; }
+        .upload-area { border: 2px dashed #ccc; padding: 40px; margin: 20px 0; border-radius: 10px; }
         .result { margin-top: 20px; }
+        .comparison { display: flex; justify-content: space-around; flex-wrap: wrap; }
+        .image-container { margin: 10px; }
+        img { max-width: 400px; height: auto; border: 1px solid #ddd; border-radius: 5px; }
+        .status { color: green; font-weight: bold; margin: 10px 0; }
+        .error { color: red; }
     </style>
 </head>
 <body>
     <div class="container">
+        <h1>🌟 ZeroIG: Zero-Shot Low-Light Enhancement</h1>
+        <p><strong>CVPR 2024</strong> - Upload a low-light image for professional enhancement!</p>
         <form method="post" enctype="multipart/form-data">
             <div class="upload-area">
                 <input type="file" name="image" accept="image/*" required>
                 <br><br>
+                <button type="submit" style="padding: 10px 20px; font-size: 16px;">🚀 Enhance with ZeroIG</button>
             </div>
         </form>
+        {% if status %}
+        <div class="status">{{ status }}</div>
+        {% endif %}
+        {% if error %}
+        <div class="error">{{ error }}</div>
+        {% endif %}
+        {% if original_image and result_image %}
         <div class="result">
+            <h3>Results:</h3>
+            <div class="comparison">
+                <div class="image-container">
+                    <h4>Original (Low-light)</h4>
+                    <img src="data:image/png;base64,{{ original_image }}" alt="Original">
+                </div>
+                <div class="image-container">
+                    <h4>ZeroIG Enhanced</h4>
+                    <img src="data:image/png;base64,{{ result_image }}" alt="Enhanced">
+                    <br><br>
+                    <a href="data:image/png;base64,{{ result_image }}" download="zeroig_enhanced.png"
+                       style="background: #007bff; color: white; padding: 10px 20px; text-decoration: none; border-radius: 5px;">
+                       📥 Download Enhanced Image
+                    </a>
+                </div>
+            </div>
         </div>
         {% endif %}
+        <div style="margin-top: 40px; padding: 20px; background: #f8f9fa; border-radius: 10px;">
+            <h3>About ZeroIG</h3>
+            <p>Zero-shot illumination-guided joint denoising and adaptive enhancement for low-light images.</p>
+            <p><strong>Features:</strong> No training data required • Joint denoising & enhancement • Prevents over-enhancement</p>
+            <p>
+                📄 <a href="https://openaccess.thecvf.com/content/CVPR2024/papers/Shi_ZERO-IG_Zero-Shot_Illumination-Guided_Joint_Denoising_and_Adaptive_Enhancement_for_Low-Light_CVPR_2024_paper.pdf">Research Paper</a> |
+                💻 <a href="https://github.com/Doyle59217/ZeroIG">Source Code</a>
+            </p>
+        </div>
     </div>
 </body>
 </html>
 """
+class ZeroIGProcessor:
+    def __init__(self):
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.model = self.load_model()
+        print(f"ZeroIG initialized on {self.device}")
+    def load_model(self):
+        try:
+            from model import Finetunemodel, Network
+            # Try to load trained weights
+            model_path = "./weights/model.pt"
+            if os.path.exists(model_path):
+                model = Finetunemodel(model_path)
+                print("✅ Loaded trained ZeroIG model")
+            else:
+                model = Network()
+                print("⚠️ Using ZeroIG with random weights (train a model for better results)")
+            model.to(self.device)
+            model.eval()
+            return model
+        except Exception as e:
+            print(f"❌ Could not load ZeroIG model: {e}")
+            return None
+    def enhance_image(self, image):
+        """Enhance image using ZeroIG model"""
+        try:
+            if self.model is None:
+                return self.simple_enhance(image), "Using simple enhancement (ZeroIG model not available)"
+            # Resize if too large to prevent memory issues
+            original_size = image.size
+            max_size = 1024
+            if max(image.size) > max_size:
+                ratio = max_size / max(image.size)
+                new_size = tuple(int(dim * ratio) for dim in image.size)
+                image = image.resize(new_size, Image.Resampling.LANCZOS)
+            # Convert to tensor
+            transform = transforms.ToTensor()
+            input_tensor = transform(image).unsqueeze(0).to(self.device)
+            # Run ZeroIG model
+            with torch.no_grad():
+                if hasattr(self.model, 'enhance') and hasattr(self.model, 'denoise_1'):
+                    # Finetunemodel - returns (enhanced, denoised)
+                    enhanced, denoised = self.model(input_tensor)
+                    result_tensor = denoised
+                    status = "✅ Enhanced with ZeroIG Finetuned model"
+                else:
+                    # Network model - returns multiple outputs
+                    outputs = self.model(input_tensor)
+                    result_tensor = outputs[13]  # H3 is the final denoised result
+                    status = "✅ Enhanced with ZeroIG Network model"
+            # Convert back to PIL
+            result_tensor = result_tensor.squeeze(0).cpu().clamp(0, 1)
+            enhanced_image = transforms.ToPILImage()(result_tensor)
+            # Resize back to original size if needed
+            if enhanced_image.size != original_size:
+                enhanced_image = enhanced_image.resize(original_size, Image.Resampling.LANCZOS)
+            return enhanced_image, status
+        except Exception as e:
+            print(f"ZeroIG enhancement error: {e}")
+            return self.simple_enhance(image), f"⚠️ ZeroIG failed, using simple enhancement: {str(e)}"
+    def simple_enhance(self, image):
+        """Fallback simple enhancement"""
+        arr = np.array(image).astype(np.float32)
+        enhanced = np.clip(arr * 1.8, 0, 255).astype(np.uint8)
         return Image.fromarray(enhanced)
+# Initialize ZeroIG processor
+print("🚀 Loading ZeroIG processor...")
+zeroig = ZeroIGProcessor()
+def image_to_base64(image):
+    """Convert PIL image to base64 string"""
+    img_buffer = io.BytesIO()
+    image.save(img_buffer, format='PNG')
+    img_str = base64.b64encode(img_buffer.getvalue()).decode()
+    return img_str
 @app.route('/', methods=['GET', 'POST'])
 def index():
+    original_image = None
     result_image = None
+    status = None
+    error = None
     if request.method == 'POST':
         try:
             file = request.files['image']
             if file:
+                print(f"Processing uploaded image: {file.filename}")
                 # Open and process image
                 image = Image.open(file.stream).convert('RGB')
+                print(f"Image size: {image.size}")
+                # Store original for comparison
+                original_image = image_to_base64(image)
+                # Enhance with ZeroIG
+                enhanced_image, enhancement_status = zeroig.enhance_image(image)
+                # Convert result to base64
+                result_image = image_to_base64(enhanced_image)
+                status = enhancement_status
         except Exception as e:
+            error = f"Error processing image: {str(e)}"
+            print(f"Error: {e}")
+    return render_template_string(HTML_TEMPLATE,
+                                  original_image=original_image,
+                                  result_image=result_image,
+                                  status=status,
+                                  error=error)
 if __name__ == '__main__':
+    print("🚀 Starting ZeroIG Flask app...")
     app.run(host='0.0.0.0', port=7860)