Spaces:

syedaoon
/

ZeroIG

Sleeping

App Files Files Community

syedaoon commited on 26 days ago

Commit

c8308c5

verified ·

1 Parent(s): d66ec3a

Update app.py

Browse files

Files changed (1) hide show

app.py +142 -169

app.py CHANGED Viewed

@@ -14,10 +14,9 @@ HTML_TEMPLATE = """
 <!DOCTYPE html>
 <html>
 <head>
-    <title>ZeroIG Enhancement - Debug</title>
     <style>
         body { font-family: Arial, sans-serif; max-width: 1000px; margin: 0 auto; padding: 20px; }
-        .debug { background: #f0f0f0; padding: 15px; margin: 10px 0; border-radius: 5px; font-family: monospace; }
         .container { text-align: center; }
         .upload-area { border: 2px dashed #ccc; padding: 40px; margin: 20px 0; border-radius: 10px; }
         .result { margin-top: 20px; }
@@ -30,23 +29,14 @@ HTML_TEMPLATE = """
 </head>
 <body>
     <div class="container">
-        <h1>🔧 ZeroIG Debug Mode</h1>
-        <div class="debug">
-            <h3>📁 File Check:</h3>
-            {{ file_status }}
-        </div>
-        <div class="debug">
-            <h3>🐍 Import Status:</h3>
-            {{ import_status }}
-        </div>
         <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;">🧪 Test Enhancement</button>
             </div>
         </form>
@@ -63,145 +53,150 @@ HTML_TEMPLATE = """
             <h3>Results:</h3>
             <div class="comparison">
                 <div class="image-container">
-                    <h4>Original</h4>
                     <img src="data:image/png;base64,{{ original_image }}" alt="Original">
                 </div>
                 <div class="image-container">
-                    <h4>Enhanced</h4>
                     <img src="data:image/png;base64,{{ result_image }}" alt="Enhanced">
                 </div>
             </div>
         </div>
         {% endif %}
     </div>
 </body>
 </html>
 """
-def check_files():
-    """Check what files are available"""
-    status = []
-    # Check current directory files
-    files_in_dir = os.listdir('.')
-    status.append(f"Files in current directory: {files_in_dir}")
-    # Check specific ZeroIG files
-    required_files = ['model.py', 'loss.py', 'utils.py']
-    for file in required_files:
-        if os.path.exists(file):
-            status.append(f"✅ {file} - EXISTS")
-            # Check file size
-            size = os.path.getsize(file)
-            status.append(f"   Size: {size} bytes")
-        else:
-            status.append(f"❌ {file} - MISSING")
-    # Check weights directory
-    if os.path.exists('weights'):
-        weights_files = os.listdir('weights')
-        status.append(f"📁 weights/ directory: {weights_files}")
-    else:
-        status.append("📁 weights/ directory: NOT FOUND")
-    return "<br>".join(status)
-def check_imports():
-    """Try importing ZeroIG modules and report results"""
-    status = []
-    try:
-        # Add current directory to path
-        import sys
-        sys.path.insert(0, '.')
-        status.append("✅ Added current directory to Python path")
-        # Try importing each module
         try:
-            import model
-            status.append("✅ Successfully imported 'model' module")
-            # Check what's in the model module
-            model_contents = dir(model)
-            status.append(f"   model module contents: {model_contents}")
-            # Try to access specific classes
-            if hasattr(model, 'Network'):
-                status.append("✅ Found Network class")
-            else:
-                status.append("❌ Network class not found")
-            if hasattr(model, 'Finetunemodel'):
-                status.append("✅ Found Finetunemodel class")
-            else:
-                status.append("❌ Finetunemodel class not found")
-        except Exception as e:
-            status.append(f"❌ Failed to import model: {e}")
-        try:
-            import loss
-            status.append("✅ Successfully imported 'loss' module")
-        except Exception as e:
-            status.append(f"❌ Failed to import loss: {e}")
-        try:
-            import utils
-            status.append("✅ Successfully imported 'utils' module")
-        except Exception as e:
-            status.append(f"❌ Failed to import utils: {e}")
-    except Exception as e:
-        status.append(f"❌ General import error: {e}")
-    return "<br>".join(status)
-def try_create_model():
-    """Try to create ZeroIG model and report what happens"""
-    try:
-        sys.path.insert(0, '.')
-        from model import Network, Finetunemodel
-        # Try creating Network
-        network = Network()
-        print("✅ Successfully created Network model")
-        # Try creating Finetunemodel (this will fail without weights, but should show the error)
         try:
-            finetuned = Finetunemodel("./weights/model.pt")
-            print("✅ Successfully created Finetunemodel")
-            return finetuned, "Finetunemodel"
-        except Exception as e:
-            print(f"⚠️ Finetunemodel failed (expected without weights): {e}")
-            print("✅ Using Network model instead")
-            return network, "Network"
-    except Exception as e:
-        print(f"❌ Model creation failed: {e}")
-        import traceback
-        traceback.print_exc()
-        return None, f"Failed: {e}"
-# Try to load model on startup
-print("🚀 Starting ZeroIG debug app...")
-print("📁 Checking files...")
-file_check = check_files()
-print("🐍 Checking imports...")
-import_check = check_imports()
-# Try to create model
-print("🤖 Trying to create model...")
-model, model_status = try_create_model()
-def simple_enhance(image):
-    """Fallback enhancement"""
-    arr = np.array(image).astype(np.float32)
-    enhanced = np.clip(arr * 1.5, 0, 255).astype(np.uint8)
-    return Image.fromarray(enhanced)
 def image_to_base64(image):
-    """Convert PIL image to base64"""
     img_buffer = io.BytesIO()
     image.save(img_buffer, format='PNG')
     img_str = base64.b64encode(img_buffer.getvalue()).decode()
@@ -218,56 +213,34 @@ def index():
         try:
             file = request.files['image']
             if file:
-                print(f"Processing: {file.filename}")
                 image = Image.open(file.stream).convert('RGB')
                 original_image = image_to_base64(image)
-                if model is not None:
-                    # Try using the actual model
-                    try:
-                        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-                        model.to(device)
-                        model.eval()
-                        transform = transforms.ToTensor()
-                        input_tensor = transform(image).unsqueeze(0).to(device)
-                        with torch.no_grad():
-                            if hasattr(model, 'enhance') and hasattr(model, 'denoise_1'):
-                                enhanced, denoised = model(input_tensor)
-                                result_tensor = denoised
-                            else:
-                                outputs = model(input_tensor)
-                                result_tensor = outputs[13]
-                        result_tensor = result_tensor.squeeze(0).cpu().clamp(0, 1)
-                        enhanced_image = transforms.ToPILImage()(result_tensor)
-                        result_image = image_to_base64(enhanced_image)
-                        status = f"✅ Used {model_status} model successfully!"
-                    except Exception as e:
-                        print(f"Model processing error: {e}")
-                        enhanced_image = simple_enhance(image)
-                        result_image = image_to_base64(enhanced_image)
-                        status = f"⚠️ Model failed, used simple enhancement: {e}"
-                else:
-                    enhanced_image = simple_enhance(image)
-                    result_image = image_to_base64(enhanced_image)
-                    status = "⚠️ No model available, used simple enhancement"
         except Exception as e:
-            error = f"Error: {e}"
             print(f"Error: {e}")
-    return render_template_string(HTML_TEMPLATE,
-                                  file_status=file_check,
-                                  import_status=import_check,
                                   original_image=original_image,
-                                  result_image=result_image,
-                                  status=status,
                                   error=error)
 if __name__ == '__main__':
-    print("🚀 Debug app ready!")
     app.run(host='0.0.0.0', port=7860)

 <!DOCTYPE html>
 <html>
 <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; }
 </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>
             <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:
+            # Import your uploaded ZeroIG files
+            from model import Finetunemodel, Network
+            # Try to load trained weights - check multiple possible paths
+            possible_weights = [
+                "./weights/LOL.pt",      # Your trained model
+                "./weights/model.pt",    # Generic name
+                "./LOL.pt",              # If uploaded to root
+                "./model.pt"             # If uploaded to root
+            ]
+            loaded_model = None
+            for weight_path in possible_weights:
+                if os.path.exists(weight_path):
+                    print(f"Found model weights at {weight_path}")
+                    try:
+                        loaded_model = Finetunemodel(weight_path)
+                        print(f"✅ Loaded ZeroIG Finetunemodel from {weight_path}")
+                        break
+                    except Exception as e:
+                        print(f"Failed to load {weight_path}: {e}")
+                        continue
+            if loaded_model is None:
+                print("No trained weights found, using Network with random initialization")
+                loaded_model = Network()
+                print("⚠️ Using ZeroIG Network with random weights - results will be poor")
+            loaded_model.to(self.device)
+            loaded_model.eval()
+            print(f"Model moved to {self.device}")
+            return loaded_model
+        except ImportError as e:
+            print(f"❌ Could not import ZeroIG modules: {e}")
+            print("Make sure you have uploaded: model.py, loss.py, utils.py")
+            return None
+        except Exception as e:
+            print(f"❌ Could not load ZeroIG model: {e}")
+            return None
+    def enhance_image(self, image):
+        """Enhance image using your ZeroIG model"""
         try:
+            if self.model is None:
+                return self.simple_enhance(image), "❌ ZeroIG model not available - using simple enhancement"
+            # Resize if too large to prevent memory issues
+            original_size = image.size
+            max_size = 800  # Adjust based on your needs
+            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)
+                print(f"Resized image from {original_size} to {image.size}")
+            # Convert to tensor
+            transform = transforms.ToTensor()
+            input_tensor = transform(image).unsqueeze(0).to(self.device)
+            print(f"Input tensor shape: {input_tensor.shape}")
+            # Run your 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  # Use denoised output
+                    status = "✅ Enhanced with ZeroIG Finetunemodel"
+                    print("Used Finetunemodel")
+                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"
+                    print("Used Network model")
+            # Convert back to PIL
+            result_tensor = result_tensor.squeeze(0).cpu().clamp(0, 1)
+            enhanced_image = transforms.ToPILImage()(result_tensor)
+            print(f"Output image size: {enhanced_image.size}")
+            # Resize back to original size if needed
+            if enhanced_image.size != original_size and original_size != image.size:
+                enhanced_image = enhanced_image.resize(original_size, Image.Resampling.LANCZOS)
+                print(f"Resized back to original size: {enhanced_image.size}")
+            return enhanced_image, status
+        except Exception as e:
+            print(f"ZeroIG enhancement error: {e}")
+            import traceback
+            traceback.print_exc()
+            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()
         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 your ZeroIG model
+                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}")
+            import traceback
+            traceback.print_exc()
+    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)