Spaces:

syedaoon
/

ZeroIG

Running

App Files Files Community

syedaoon commited on Jul 27

Commit

eae91d6

verified ·

1 Parent(s): 3b2db10

Update app.py

Browse files

Files changed (1) hide show

app.py +169 -128

app.py CHANGED Viewed

@@ -14,9 +14,10 @@ HTML_TEMPLATE = """
 <!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; }
@@ -29,14 +30,23 @@ HTML_TEMPLATE = """
 </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>
@@ -53,136 +63,145 @@ HTML_TEMPLATE = """
             <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
-            model_path = "./weights/model.pt"
-            if os.path.exists(model_path):
-                print(f"Found model weights at {model_path}")
-                model = Finetunemodel(model_path)
-                print("✅ Loaded ZeroIG Finetunemodel with trained weights")
-            else:
-                print("No trained weights found, using Network with random initialization")
-                model = Network()
-                print("⚠️ Using ZeroIG Network with random weights")
-            model.to(self.device)
-            model.eval()
-            print(f"Model moved to {self.device}")
-            return 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()
@@ -199,34 +218,56 @@ def index():
         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)

 <!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; }
 </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>
             <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()
         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)