216 lines
7.1 KiB
Python
216 lines
7.1 KiB
Python
from flask import Flask, render_template, request, redirect, url_for, send_from_directory, session, jsonify
|
|
from flask_socketio import SocketIO, emit, disconnect
|
|
import os
|
|
import json
|
|
from werkzeug.security import check_password_hash
|
|
from PIL import Image, UnidentifiedImageError
|
|
from watchdog.observers import Observer
|
|
from watchdog.events import FileSystemEventHandler
|
|
from os.path import realpath, join
|
|
|
|
app = Flask(__name__)
|
|
|
|
# Load configuration from keys.json
|
|
with open('db/keys.json', 'r') as f:
|
|
config = json.load(f)
|
|
|
|
app.config['SECRET_KEY'] = config['key']
|
|
socketio = SocketIO(app)
|
|
|
|
app.config['UPLOAD_FOLDER'] = '/home/kenny/workspace/Image_Gen_Server/output/expo_postprocessed' # Folder to store images
|
|
app.config['DUMP_FOLDER'] = '/home/kenny/workspace/ComfyUI-halloween/output' # Folder to store dump images
|
|
app.config['CACHE_FOLDER'] = 'cache' # Folder to store cached resized images
|
|
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg', 'gif'}
|
|
|
|
USERNAME = config['login_info']['username']
|
|
PASSWORD_HASH = config['login_info']['password_hash']
|
|
SALT = config['login_info']['salt']
|
|
|
|
def allowed_file(filename):
|
|
return '.' in filename and \
|
|
filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']
|
|
|
|
def is_image_valid(image_path):
|
|
try:
|
|
with Image.open(image_path) as img:
|
|
img.verify() # Verify that it is, in fact, an image
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
def resize_image(image_path, cache_path, size=(800, 450)):
|
|
with Image.open(image_path) as img:
|
|
img.thumbnail(size)
|
|
if img.mode in ("RGBA", "P"): # Convert to RGB if necessary
|
|
img = img.convert("RGB")
|
|
img.save(cache_path, format="JPEG")
|
|
|
|
def create_cache():
|
|
cached_images = []
|
|
|
|
if not os.path.exists(app.config['CACHE_FOLDER']):
|
|
os.makedirs(app.config['CACHE_FOLDER'])
|
|
|
|
for image in os.listdir(app.config['UPLOAD_FOLDER']):
|
|
image_path = os.path.join(app.config['UPLOAD_FOLDER'], image)
|
|
cache_path = os.path.join(app.config['CACHE_FOLDER'], image)
|
|
|
|
if not os.path.exists(cache_path) and is_image_valid(image_path):
|
|
resize_image(image_path, cache_path)
|
|
|
|
if is_image_valid(image_path):
|
|
cached_images.append(image)
|
|
|
|
return cached_images
|
|
|
|
@app.route('/')
|
|
def index():
|
|
return redirect(url_for('login'))
|
|
|
|
@app.route('/login', methods=['GET', 'POST'])
|
|
def login():
|
|
if request.method == 'POST':
|
|
username = request.form['username']
|
|
password = request.form['password']
|
|
if username == USERNAME and check_password_hash(PASSWORD_HASH, password + SALT):
|
|
session['logged_in'] = True
|
|
return redirect(url_for('gallery'))
|
|
else:
|
|
error = 'Invalid credentials'
|
|
return render_template('login.html', error=error)
|
|
return render_template('login.html')
|
|
|
|
@app.route('/logout')
|
|
def logout():
|
|
session.pop('logged_in', None)
|
|
return redirect(url_for('login'))
|
|
|
|
@app.route('/gallery')
|
|
def gallery():
|
|
if not session.get('logged_in'):
|
|
return redirect(url_for('login'))
|
|
|
|
cached_images = create_cache()
|
|
batches = []
|
|
current_batch = []
|
|
current_batch_number = None
|
|
unstructured_images = []
|
|
|
|
for image in cached_images:
|
|
parts = image.split('_')
|
|
if len(parts) == 2 and parts[1].split('.')[0].isdigit():
|
|
batch_number = int(parts[0])
|
|
if current_batch_number is None:
|
|
current_batch_number = batch_number
|
|
if batch_number != current_batch_number:
|
|
batches.append(current_batch)
|
|
current_batch = []
|
|
current_batch_number = batch_number
|
|
current_batch.append(image)
|
|
else:
|
|
unstructured_images.append(image)
|
|
|
|
if current_batch:
|
|
batches.append(current_batch)
|
|
if unstructured_images:
|
|
batches.append(unstructured_images)
|
|
|
|
return render_template('gallery.html', batches=batches)
|
|
|
|
@app.route('/images/<filename>')
|
|
def uploaded_file(filename):
|
|
if not session.get('logged_in'):
|
|
return redirect(url_for('login'))
|
|
|
|
image_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
|
|
if is_image_valid(image_path):
|
|
return send_from_directory(app.config['UPLOAD_FOLDER'], filename)
|
|
else:
|
|
return "Invalid image", 404
|
|
|
|
@app.route('/cached/<filename>')
|
|
def cached_file(filename):
|
|
if not session.get('logged_in'):
|
|
return redirect(url_for('login'))
|
|
|
|
image_path = os.path.join(app.config['CACHE_FOLDER'], filename)
|
|
if is_image_valid(image_path):
|
|
return send_from_directory(app.config['CACHE_FOLDER'], filename)
|
|
else:
|
|
return "Invalid image", 404
|
|
|
|
@app.route('/dump/<filename>')
|
|
def dump_file(filename):
|
|
if not session.get('logged_in'):
|
|
return redirect(url_for('login'))
|
|
|
|
image_path = os.path.join(app.config['DUMP_FOLDER'], filename)
|
|
if is_image_valid(image_path):
|
|
return send_from_directory(app.config['DUMP_FOLDER'], filename)
|
|
else:
|
|
return "Invalid image", 404
|
|
|
|
@app.route('/api/images')
|
|
def api_images():
|
|
if not session.get('logged_in'):
|
|
return jsonify([])
|
|
|
|
images = sorted(
|
|
os.listdir(app.config['UPLOAD_FOLDER']),
|
|
key=lambda x: os.path.getctime(os.path.join(app.config['UPLOAD_FOLDER'], x)),
|
|
reverse=True
|
|
)
|
|
valid_images = [img for img in images if is_image_valid(os.path.join(app.config['UPLOAD_FOLDER'], img))]
|
|
return jsonify(valid_images)
|
|
|
|
@app.route('/dump')
|
|
def dump():
|
|
if not session.get('logged_in'):
|
|
return redirect(url_for('login'))
|
|
|
|
#cached_images = create_cache()
|
|
dump_folder = app.config['DUMP_FOLDER'] # Folder to store dump images
|
|
images = sorted(
|
|
os.listdir(dump_folder),
|
|
key=lambda x: os.path.getctime(os.path.join(dump_folder, x)),
|
|
reverse=True
|
|
)
|
|
valid_images = [img for img in images if is_image_valid(os.path.join(dump_folder, img))]
|
|
print(valid_images)
|
|
return render_template('dump.html', images=valid_images)
|
|
|
|
def emit_gallery_update():
|
|
socketio.emit('update_gallery')
|
|
|
|
class Watcher(FileSystemEventHandler):
|
|
def on_modified(self, event):
|
|
if not event.is_directory and event.src_path.startswith(app.config['UPLOAD_FOLDER']):
|
|
create_cache()
|
|
emit_gallery_update()
|
|
|
|
def on_created(self, event):
|
|
if not event.is_directory and event.src_path.startswith(app.config['UPLOAD_FOLDER']):
|
|
create_cache()
|
|
emit_gallery_update()
|
|
|
|
def on_deleted(self, event):
|
|
if not event.is_directory and event.src_path.startswith(app.config['UPLOAD_FOLDER']):
|
|
create_cache()
|
|
emit_gallery_update()
|
|
|
|
@socketio.on('connect')
|
|
def handle_connect():
|
|
if not session.get('logged_in'):
|
|
disconnect()
|
|
|
|
if __name__ == '__main__':
|
|
observer = Observer()
|
|
observer.schedule(Watcher(), path=app.config['UPLOAD_FOLDER'], recursive=False)
|
|
observer.start()
|
|
try:
|
|
context = ('cert.pem', 'key.pem') # Replace with your self-signed certificate and key files
|
|
socketio.run(app, debug=True, ssl_context=context)
|
|
except KeyboardInterrupt:
|
|
observer.stop()
|
|
observer.join()
|