|
|
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 |
|
|
|
|
|
app = Flask(__name__) |
|
|
|
|
|
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; } |
|
|
.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 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> |
|
|
|
|
|
{% 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</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 = [] |
|
|
|
|
|
|
|
|
files_in_dir = os.listdir('.') |
|
|
status.append(f"Files in current directory: {files_in_dir}") |
|
|
|
|
|
|
|
|
required_files = ['model.py', 'loss.py', 'utils.py'] |
|
|
for file in required_files: |
|
|
if os.path.exists(file): |
|
|
status.append(f"✅ {file} - EXISTS") |
|
|
|
|
|
size = os.path.getsize(file) |
|
|
status.append(f" Size: {size} bytes") |
|
|
else: |
|
|
status.append(f"❌ {file} - MISSING") |
|
|
|
|
|
|
|
|
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: |
|
|
|
|
|
import sys |
|
|
sys.path.insert(0, '.') |
|
|
status.append("✅ Added current directory to Python path") |
|
|
|
|
|
|
|
|
try: |
|
|
import model |
|
|
status.append("✅ Successfully imported 'model' module") |
|
|
|
|
|
|
|
|
model_contents = dir(model) |
|
|
status.append(f" model module contents: {model_contents}") |
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
network = Network() |
|
|
print("✅ Successfully created Network model") |
|
|
|
|
|
|
|
|
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}" |
|
|
|
|
|
|
|
|
print("🚀 Starting ZeroIG debug app...") |
|
|
print("📁 Checking files...") |
|
|
file_check = check_files() |
|
|
print("🐍 Checking imports...") |
|
|
import_check = check_imports() |
|
|
|
|
|
|
|
|
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() |
|
|
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: {file.filename}") |
|
|
|
|
|
image = Image.open(file.stream).convert('RGB') |
|
|
original_image = image_to_base64(image) |
|
|
|
|
|
if model is not None: |
|
|
|
|
|
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) |
