File size: 5,668 Bytes
c33ad35 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
import os
import uuid
from flask import Flask, render_template, request, jsonify, send_from_directory, url_for
from werkzeug.utils import secure_filename
import cv2
import numpy as np
from PIL import Image
from image_to_sketch import ImageToSketch
app = Flask(__name__)
# Configuration
app.config['MAX_CONTENT_LENGTH'] = 10 * 1024 * 1024 # 10MB max file size
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['RESULT_FOLDER'] = 'results'
# Allowed file extensions
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'tiff', 'webp'}
def allowed_file(filename):
"""Check if file extension is allowed."""
return '.' in filename and \
filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
def ensure_directories():
"""Ensure upload and result directories exist."""
for directory in [app.config['UPLOAD_FOLDER'], app.config['RESULT_FOLDER']]:
if not os.path.exists(directory):
os.makedirs(directory)
@app.route('/')
def index():
"""Serve the main page."""
return render_template('index.html')
@app.route('/convert', methods=['POST'])
def convert_image():
"""Handle image conversion to sketch."""
try:
# Check if image was uploaded
if 'image' not in request.files:
return jsonify({'error': 'No image file provided'}), 400
file = request.files['image']
if file.filename == '':
return jsonify({'error': 'No file selected'}), 400
if not allowed_file(file.filename):
return jsonify({'error': 'Invalid file type. Please upload an image file.'}), 400
# Get processing options
blur_value = int(request.form.get('blur', 21))
enhance = request.form.get('enhance', 'false').lower() == 'true'
# Validate blur value
if blur_value < 1:
return jsonify({'error': 'Blur value must be positive'}), 400
# Make blur value odd
if blur_value % 2 == 0:
blur_value += 1
# Generate unique filename
file_id = str(uuid.uuid4())
if file.filename and '.' in file.filename:
original_ext = file.filename.rsplit('.', 1)[1].lower()
else:
original_ext = 'jpg' # Default extension
input_filename = f"{file_id}_original.{original_ext}"
output_filename = f"{file_id}_sketch.jpg"
# Save uploaded file
input_path = os.path.join(app.config['UPLOAD_FOLDER'], input_filename)
output_path = os.path.join(app.config['RESULT_FOLDER'], output_filename)
file.save(input_path)
# Convert image to sketch
converter = ImageToSketch()
sketch = converter.convert_to_sketch(
input_path,
output_path,
blur_value=blur_value,
enhance=enhance
)
# Generate URL for the result
sketch_url = url_for('get_result', filename=output_filename)
# Clean up input file after successful processing
try:
os.remove(input_path)
except Exception:
pass # Don't fail if cleanup fails
return jsonify({
'success': True,
'sketch_url': sketch_url,
'blur_value': blur_value,
'enhanced': enhance
})
except cv2.error as e:
print(f"OpenCV error processing image: {str(e)}")
return jsonify({'error': 'Invalid or corrupted image file'}), 400
except Exception as e:
print(f"Error processing image: {str(e)}")
return jsonify({'error': 'Failed to process image. Please try again.'}), 500
@app.route('/results/<filename>')
def get_result(filename):
"""Serve result images."""
return send_from_directory(app.config['RESULT_FOLDER'], filename)
# Removed /uploads route for security - originals should not be publicly accessible
@app.errorhandler(413)
def too_large(e):
"""Handle file too large error."""
return jsonify({'error': 'File size exceeds 10MB limit'}), 413
@app.errorhandler(500)
def internal_error(e):
"""Handle internal server errors."""
return jsonify({'error': 'Internal server error'}), 500
def cleanup_old_files():
"""Clean up old uploaded and result files."""
import time
current_time = time.time()
max_age = 24 * 60 * 60 # 24 hours
for folder in [app.config['UPLOAD_FOLDER'], app.config['RESULT_FOLDER']]:
if os.path.exists(folder):
for filename in os.listdir(folder):
file_path = os.path.join(folder, filename)
if os.path.isfile(file_path):
file_age = current_time - os.path.getctime(file_path)
if file_age > max_age:
try:
os.remove(file_path)
print(f"Cleaned up old file: {file_path}")
except Exception as e:
print(f"Error cleaning up {file_path}: {e}")
if __name__ == '__main__':
ensure_directories()
# Clean up old files on startup
try:
cleanup_old_files()
except Exception as e:
print(f"Error during cleanup: {e}")
# Auto-detect port: prefer PORT env var, else 7860 for HF, else 5000 for local
port = int(os.environ.get('PORT') or (7860 if 'SPACE_ID' in os.environ else 5000))
# Start the Flask development server
print("Starting SketchAI Web Application...")
print(f"Open your browser and navigate to: http://localhost:{port}")
app.run(host='0.0.0.0', port=port, debug=False) |