|
|
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__) |
|
|
|
|
|
|
|
|
sys.path.append('.') |
|
|
sys.path.append('./CFWD') |
|
|
|
|
|
HTML_TEMPLATE = """ |
|
|
<!DOCTYPE html> |
|
|
<html> |
|
|
<head> |
|
|
<title>Multi-Method Low-Light Enhancement</title> |
|
|
<style> |
|
|
body { font-family: Arial, sans-serif; max-width: 1200px; margin: 0 auto; padding: 20px; } |
|
|
.container { text-align: center; } |
|
|
.method-selection { margin: 20px 0; padding: 20px; background: #f8f9fa; border-radius: 10px; } |
|
|
.method-button { |
|
|
margin: 5px; padding: 10px 20px; border: none; border-radius: 5px; cursor: pointer; |
|
|
font-size: 14px; transition: all 0.3s; |
|
|
} |
|
|
.method-button.active { background: #007bff; color: white; } |
|
|
.method-button.inactive { background: #e9ecef; color: #6c757d; } |
|
|
.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; } |
|
|
.image-container img { max-width: 350px; height: auto; border: 1px solid #ddd; border-radius: 5px; } |
|
|
.status { color: green; font-weight: bold; margin: 10px 0; } |
|
|
.error { color: red; } |
|
|
.processing { color: orange; font-weight: bold; } |
|
|
.method-info { background: #e3f2fd; padding: 15px; margin: 10px 0; border-radius: 5px; } |
|
|
.folder-status { background: #fff3cd; padding: 10px; margin: 10px 0; border-radius: 5px; font-family: monospace; } |
|
|
</style> |
|
|
<script> |
|
|
function selectMethod(method) { |
|
|
document.getElementById('selected_method').value = method; |
|
|
|
|
|
// Update button styles |
|
|
const buttons = document.querySelectorAll('.method-button'); |
|
|
buttons.forEach(btn => { |
|
|
if (btn.dataset.method === method) { |
|
|
btn.classList.add('active'); |
|
|
btn.classList.remove('inactive'); |
|
|
} else { |
|
|
btn.classList.add('inactive'); |
|
|
btn.classList.remove('active'); |
|
|
} |
|
|
}); |
|
|
|
|
|
// Update method info |
|
|
const infos = document.querySelectorAll('.method-info'); |
|
|
infos.forEach(info => { |
|
|
info.style.display = info.dataset.method === method ? 'block' : 'none'; |
|
|
}); |
|
|
} |
|
|
|
|
|
window.onload = function() { |
|
|
selectMethod('zeroig'); |
|
|
} |
|
|
</script> |
|
|
</head> |
|
|
<body> |
|
|
<div class="container"> |
|
|
<h1>🌟 Multi-Method Low-Light Enhancement</h1> |
|
|
<p>Professional low-light image enhancement with multiple state-of-the-art methods</p> |
|
|
|
|
|
<div class="folder-status"> |
|
|
<strong>📁 Method Status:</strong><br> |
|
|
{{ method_status }} |
|
|
</div> |
|
|
|
|
|
<div class="method-selection"> |
|
|
<h3>📚 Select Enhancement Method:</h3> |
|
|
<button class="method-button" data-method="zeroig" onclick="selectMethod('zeroig')"> |
|
|
🎯 ZeroIG (CVPR 2024) |
|
|
</button> |
|
|
<button class="method-button" data-method="cfwd" onclick="selectMethod('cfwd')"> |
|
|
🌊 CFWD (CLIP-Fourier Wavelet) |
|
|
</button> |
|
|
<button class="method-button" data-method="both" onclick="selectMethod('both')"> |
|
|
🔄 Compare Both Methods |
|
|
</button> |
|
|
|
|
|
<div class="method-info" data-method="zeroig"> |
|
|
<strong>🎯 ZeroIG:</strong> Zero-shot illumination-guided joint denoising and adaptive enhancement. |
|
|
Fast processing, no training data required, prevents over-enhancement artifacts. |
|
|
</div> |
|
|
|
|
|
<div class="method-info" data-method="cfwd"> |
|
|
<strong>🌊 CFWD:</strong> CLIP-Fourier Guided Wavelet Diffusion for low-light enhancement. |
|
|
Uses diffusion models with CLIP guidance for state-of-the-art quality. |
|
|
</div> |
|
|
|
|
|
<div class="method-info" data-method="both"> |
|
|
<strong>🔄 Comparison Mode:</strong> Process with both methods to see the differences. |
|
|
Takes longer but shows you the best of both approaches side-by-side. |
|
|
</div> |
|
|
</div> |
|
|
|
|
|
<form method="post" enctype="multipart/form-data"> |
|
|
<input type="hidden" id="selected_method" name="method" value="zeroig"> |
|
|
|
|
|
<div class="upload-area"> |
|
|
<input type="file" name="image" accept="image/*" required> |
|
|
<br><br> |
|
|
<button type="submit" style="padding: 15px 30px; font-size: 16px; background: #28a745; color: white; border: none; border-radius: 5px;"> |
|
|
🚀 Enhance Image |
|
|
</button> |
|
|
</div> |
|
|
</form> |
|
|
|
|
|
{% if status %} |
|
|
<div class="status">{{ status }}</div> |
|
|
{% endif %} |
|
|
|
|
|
{% if error %} |
|
|
<div class="error">{{ error }}</div> |
|
|
{% endif %} |
|
|
|
|
|
{% if results %} |
|
|
<div class="result"> |
|
|
<h3>🎯 Enhancement Results:</h3> |
|
|
<div class="comparison"> |
|
|
<div class="image-container"> |
|
|
<h4>📷 Original (Low-light)</h4> |
|
|
<img src="data:image/png;base64,{{ results.original }}" alt="Original"> |
|
|
</div> |
|
|
|
|
|
{% if results.zeroig %} |
|
|
<div class="image-container"> |
|
|
<h4>🎯 ZeroIG Enhanced</h4> |
|
|
<img src="data:image/png;base64,{{ results.zeroig }}" alt="ZeroIG"> |
|
|
<br><br> |
|
|
<a href="data:image/png;base64,{{ results.zeroig }}" download="zeroig_enhanced.png" |
|
|
style="background: #007bff; color: white; padding: 8px 16px; text-decoration: none; border-radius: 3px;"> |
|
|
📥 Download ZeroIG |
|
|
</a> |
|
|
</div> |
|
|
{% endif %} |
|
|
|
|
|
{% if results.cfwd %} |
|
|
<div class="image-container"> |
|
|
<h4>🌊 CFWD Enhanced</h4> |
|
|
<img src="data:image/png;base64,{{ results.cfwd }}" alt="CFWD"> |
|
|
<br><br> |
|
|
<a href="data:image/png;base64,{{ results.cfwd }}" download="cfwd_enhanced.png" |
|
|
style="background: #17a2b8; color: white; padding: 8px 16px; text-decoration: none; border-radius: 3px;"> |
|
|
📥 Download CFWD |
|
|
</a> |
|
|
</div> |
|
|
{% endif %} |
|
|
</div> |
|
|
|
|
|
{% if method_used %} |
|
|
<div style="margin-top: 20px; padding: 15px; background: #d4edda; border-radius: 5px;"> |
|
|
<strong>Method Used:</strong> {{ method_used }}<br> |
|
|
<strong>Processing Time:</strong> {{ processing_time }}s |
|
|
</div> |
|
|
{% endif %} |
|
|
</div> |
|
|
{% endif %} |
|
|
|
|
|
<div style="margin-top: 40px; padding: 20px; background: #f8f9fa; border-radius: 10px;"> |
|
|
<h3>📖 About the Methods</h3> |
|
|
|
|
|
<div style="margin: 15px 0;"> |
|
|
<strong>🎯 ZeroIG (CVPR 2024):</strong> |
|
|
<p>Zero-shot illumination-guided joint denoising and adaptive enhancement for low-light images.</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">Paper</a> | |
|
|
💻 <a href="https://github.com/Doyle59217/ZeroIG">Code</a></p> |
|
|
</div> |
|
|
|
|
|
<div style="margin: 15px 0;"> |
|
|
<strong>🌊 CFWD (2024):</strong> |
|
|
<p>Low-light Image Enhancement via CLIP-Fourier Guided Wavelet Diffusion.</p> |
|
|
<p>📄 <a href="https://arxiv.org/abs/2401.03788">Paper</a> | |
|
|
💻 <a href="https://github.com/hejh8/CFWD">Code</a></p> |
|
|
</div> |
|
|
</div> |
|
|
</div> |
|
|
</body> |
|
|
</html> |
|
|
""" |
|
|
|
|
|
def check_method_availability(): |
|
|
"""Check which methods are available""" |
|
|
status = [] |
|
|
|
|
|
|
|
|
zeroig_files = ['model.py', 'loss.py', 'utils.py'] |
|
|
zeroig_ok = all(os.path.exists(f) for f in zeroig_files) |
|
|
|
|
|
if zeroig_ok: |
|
|
|
|
|
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") |
|
|
|
|
|
|
|
|
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") |
|
|
|
|
|
|
|
|
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 "<br>".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 |
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
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: |
|
|
|
|
|
if not os.path.exists('./CFWD'): |
|
|
print("❌ CFWD folder not found") |
|
|
return None |
|
|
|
|
|
|
|
|
sys.path.insert(0, './CFWD') |
|
|
|
|
|
|
|
|
|
|
|
try: |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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: |
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
transform = transforms.ToTensor() |
|
|
input_tensor = transform(image).unsqueeze(0).to(self.device) |
|
|
|
|
|
|
|
|
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] |
|
|
|
|
|
|
|
|
result_tensor = result_tensor.squeeze(0).cpu().clamp(0, 1) |
|
|
enhanced_image = transforms.ToPILImage()(result_tensor) |
|
|
|
|
|
|
|
|
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: |
|
|
|
|
|
try: |
|
|
|
|
|
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: |
|
|
|
|
|
|
|
|
pass |
|
|
except Exception as e: |
|
|
return None, f"CFWD failed: {str(e)}" |
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
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() |
|
|
|
|
|
|
|
|
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': |
|
|
|
|
|
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) |
