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 import sys import time app = Flask(__name__) # Add paths for different methods sys.path.append('.') # For ZeroIG (root folder) sys.path.append('./CFWD') # For CFWD method HTML_TEMPLATE = """ Multi-Method Low-Light Enhancement

🌟 Multi-Method Low-Light Enhancement

Professional low-light image enhancement with multiple state-of-the-art methods

📁 Method Status:
{{ method_status }}

📚 Select Enhancement Method:

🎯 ZeroIG: Zero-shot illumination-guided joint denoising and adaptive enhancement. Fast processing, no training data required, prevents over-enhancement artifacts.
🌊 CFWD: CLIP-Fourier Guided Wavelet Diffusion for low-light enhancement. Uses diffusion models with CLIP guidance for state-of-the-art quality.
🔄 Comparison Mode: Process with both methods to see the differences. Takes longer but shows you the best of both approaches side-by-side.


{% if status %}
{{ status }}
{% endif %} {% if error %}
{{ error }}
{% endif %} {% if results %}

🎯 Enhancement Results:

📷 Original (Low-light)

Original
{% if results.zeroig %}

🎯 ZeroIG Enhanced

ZeroIG

📥 Download ZeroIG
{% endif %} {% if results.cfwd %}

🌊 CFWD Enhanced

CFWD

📥 Download CFWD
{% endif %}
{% if method_used %}
Method Used: {{ method_used }}
Processing Time: {{ processing_time }}s
{% endif %}
{% endif %}

📖 About the Methods

🎯 ZeroIG (CVPR 2024):

Zero-shot illumination-guided joint denoising and adaptive enhancement for low-light images.

📄 Paper | 💻 Code

🌊 CFWD (2024):

Low-light Image Enhancement via CLIP-Fourier Guided Wavelet Diffusion.

📄 Paper | 💻 Code

""" def check_method_availability(): """Check which methods are available""" status = [] # Check ZeroIG (root folder) zeroig_files = ['model.py', 'loss.py', 'utils.py'] zeroig_ok = all(os.path.exists(f) for f in zeroig_files) if zeroig_ok: # Check for weights weights_found = False for weight_path in ["./weights/LOL.pt", "./weights/model.pt"]: if os.path.exists(weight_path): weights_found = True break if weights_found: status.append("✅ ZeroIG: Ready with trained weights") else: status.append("⚠️ ZeroIG: Available but no trained weights found") else: status.append("❌ ZeroIG: Missing required files") # Check CFWD folder cfwd_folder_exists = os.path.exists('./CFWD') if cfwd_folder_exists: cfwd_files = os.listdir('./CFWD') status.append(f"📁 CFWD folder: Found with {len(cfwd_files)} files") # Check for common CFWD files expected_cfwd_files = ['test.py', 'model.py', 'datasets.py'] found_files = [f for f in expected_cfwd_files if f in cfwd_files] if found_files: status.append(f"✅ CFWD files found: {', '.join(found_files)}") else: status.append("⚠️ CFWD: Folder exists but expected files not found") else: status.append("❌ CFWD: Folder not found - please create ./CFWD/ and upload CFWD files") return "
".join(status) class MultiMethodProcessor: def __init__(self): self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.zeroig_model = self.load_zeroig() self.cfwd_model = self.load_cfwd() print(f"Multi-method processor initialized on {self.device}") def load_zeroig(self): """Load ZeroIG model from root folder""" try: from model import Finetunemodel, Network # Try to load trained weights possible_weights = [ "./weights/LOL.pt", "./weights/zeroig.pt", "./weights/model.pt" ] for weight_path in possible_weights: if os.path.exists(weight_path): try: model = Finetunemodel(weight_path) model.to(self.device) model.eval() print(f"✅ ZeroIG loaded from {weight_path}") return model except Exception as e: print(f"Failed to load ZeroIG from {weight_path}: {e}") continue # Fallback to Network model = Network() model.to(self.device) model.eval() print("⚠️ ZeroIG using random weights") return model except Exception as e: print(f"❌ ZeroIG not available: {e}") return None def load_cfwd(self): """Load CFWD model from CFWD folder""" try: # Check if CFWD folder exists if not os.path.exists('./CFWD'): print("❌ CFWD folder not found") return None # Add CFWD folder to path and try importing sys.path.insert(0, './CFWD') # Try to import CFWD modules (will depend on actual CFWD structure) # This is a placeholder - you'll need to adapt based on actual CFWD files try: # Example imports - replace with actual CFWD imports # from model import CFWDModel # Replace with actual CFWD model # model = CFWDModel() # model.load_weights('./CFWD/weights/cfwd_weights.pt') print("⚠️ CFWD: Import structure not yet implemented") return None except ImportError as e: print(f"❌ CFWD import failed: {e}") return None except Exception as e: print(f"❌ CFWD loading error: {e}") return None def enhance_with_zeroig(self, image): """Enhance image using ZeroIG""" if self.zeroig_model is None: return None, "ZeroIG model not available" try: # Resize if needed original_size = image.size max_size = 800 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 with torch.no_grad(): if hasattr(self.zeroig_model, 'enhance') and hasattr(self.zeroig_model, 'denoise_1'): enhanced, denoised = self.zeroig_model(input_tensor) result_tensor = denoised else: outputs = self.zeroig_model(input_tensor) result_tensor = outputs[13] # Convert back to PIL result_tensor = result_tensor.squeeze(0).cpu().clamp(0, 1) enhanced_image = transforms.ToPILImage()(result_tensor) # Resize back if needed if enhanced_image.size != original_size: enhanced_image = enhanced_image.resize(original_size, Image.Resampling.LANCZOS) return enhanced_image, "ZeroIG enhancement successful" except Exception as e: return None, f"ZeroIG failed: {str(e)}" def enhance_with_cfwd(self, image): """Enhance image using CFWD""" if self.cfwd_model is None: # Placeholder enhancement until CFWD is implemented try: # Simple placeholder - replace with actual CFWD processing arr = np.array(image).astype(np.float32) enhanced = np.clip(arr * 2.2, 0, 255).astype(np.uint8) enhanced_image = Image.fromarray(enhanced) return enhanced_image, "CFWD placeholder (upload CFWD files to ./CFWD/ folder for real processing)" except Exception as e: return None, f"CFWD placeholder failed: {str(e)}" try: # TODO: Implement actual CFWD processing when files are uploaded # This will depend on the actual CFWD model structure pass except Exception as e: return None, f"CFWD failed: {str(e)}" # Initialize processor print("🚀 Loading multi-method processor...") processor = MultiMethodProcessor() 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() return img_str @app.route('/', methods=['GET', 'POST']) def index(): results = None status = None error = None method_used = None processing_time = None # Check method availability method_status = check_method_availability() if request.method == 'POST': try: file = request.files['image'] method = request.form.get('method', 'zeroig') if file: print(f"Processing with method: {method}") start_time = time.time() # Load image image = Image.open(file.stream).convert('RGB') original_b64 = image_to_base64(image) results = {'original': original_b64} if method == 'zeroig': enhanced, msg = processor.enhance_with_zeroig(image) if enhanced: results['zeroig'] = image_to_base64(enhanced) status = f"✅ ZeroIG: {msg}" else: error = f"❌ ZeroIG failed: {msg}" elif method == 'cfwd': enhanced, msg = processor.enhance_with_cfwd(image) if enhanced: results['cfwd'] = image_to_base64(enhanced) status = f"✅ CFWD: {msg}" else: error = f"❌ CFWD failed: {msg}" elif method == 'both': # Process with both methods zeroig_enhanced, zeroig_msg = processor.enhance_with_zeroig(image) cfwd_enhanced, cfwd_msg = processor.enhance_with_cfwd(image) status_msgs = [] if zeroig_enhanced: results['zeroig'] = image_to_base64(zeroig_enhanced) status_msgs.append(f"ZeroIG: {zeroig_msg}") if cfwd_enhanced: results['cfwd'] = image_to_base64(cfwd_enhanced) status_msgs.append(f"CFWD: {cfwd_msg}") status = "✅ " + " | ".join(status_msgs) end_time = time.time() processing_time = round(end_time - start_time, 2) method_used = method.upper() except Exception as e: error = f"Error processing image: {str(e)}" print(f"Error: {e}") return render_template_string(HTML_TEMPLATE, method_status=method_status, results=results, status=status, error=error, method_used=method_used, processing_time=processing_time) if __name__ == '__main__': print("🚀 Starting Organized Multi-Method Enhancement App...") app.run(host='0.0.0.0', port=7860)