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/') 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)