sd_image_selector/app.py

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

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'] = 'images'  # Folder to store images
app.config['DUMP_FOLDER'] = 'dump_images'
app.config['CACHE_FOLDER'] = 'cache'  # Folder to store cached resized images
app.config['DUMP_CACHE_FOLDER'] = 'dump_cache'
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 delete_non_existing_cache():
    cache_folder = app.config['CACHE_FOLDER']
    upload_folder = app.config['UPLOAD_FOLDER']
    dump_cache_folder = app.config['DUMP_CACHE_FOLDER']
    dump_folder = app.config['DUMP_FOLDER']

    images = os.listdir(cache_folder)
    for image in images:
        image_path_upload = os.path.join(upload_folder, image)
        cache_path = os.path.join(cache_folder, image)

        if not os.path.exists(image_path_upload):
            if os.path.exists(cache_path):
                os.remove(cache_path)
    
    dump_images = os.listdir(dump_cache_folder)
    for image in dump_images:
        image_path_dump = os.path.join(dump_folder, image)
        dump_cache_path = os.path.join(dump_cache_folder, image)

        if not os.path.exists(image_path_dump):
            if os.path.exists(dump_cache_path):
                os.remove(dump_cache_path)

def create_cache():
    images = os.listdir(app.config['UPLOAD_FOLDER'])
    cached_images = []

    if not os.path.exists(app.config['CACHE_FOLDER']):
        os.makedirs(app.config['CACHE_FOLDER'])

    caches_path = set([os.path.join(app.config['CACHE_FOLDER'], image) for image in os.listdir(app.config['CACHE_FOLDER'])])
    images_path = set([os.path.join(app.config['UPLOAD_FOLDER'], image) for image in images])

    # get the images that are not cached
    uncached_images = images_path - caches_path
    for image_path in uncached_images:
        if is_image_valid(image_path):
            cache_path = os.path.join(app.config['CACHE_FOLDER'], os.path.basename(image_path))
            resize_image(image_path, cache_path)
            cached_images.append(os.path.basename(image_path)) 
    
    return cached_images

def create_dump_cache():
    images = os.listdir(app.config['DUMP_FOLDER'])
    cached_images = []

    if not os.path.exists(app.config['DUMP_CACHE_FOLDER']):
        os.makedirs(app.config['DUMP_CACHE_FOLDER'])

    caches_path = set([os.path.join(app.config['DUMP_CACHE_FOLDER'], image) for image in os.listdir(app.config['DUMP_CACHE_FOLDER'])])
    images_path = set([os.path.join(app.config['DUMP_FOLDER'], image) for image in images])

    # get the images that are not cached
    uncached_images = images_path - caches_path
    for image_path in uncached_images:
        if is_image_valid(image_path):
            cache_path = os.path.join(app.config['DUMP_CACHE_FOLDER'], os.path.basename(image_path))
            resize_image(image_path, cache_path)
            cached_images.append(os.path.basename(image_path)) 
    
    return cached_images

def update_cache():
    create_cache()
    delete_non_existing_cache()

def update_dump_cache():
    create_dump_cache()
    delete_non_existing_cache()
    
@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()
    cached_images.sort()
    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)

    batches = batches[::-1]

    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('/dump_images/<filename>')
def uploaded_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('/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_cached/<filename>')
def dump_cached_file(filename):
    if not session.get('logged_in'):
        return redirect(url_for('login'))
    
    image_path = os.path.join(app.config['DUMP_CACHE_FOLDER'], filename)
    if is_image_valid(image_path):
        return send_from_directory(app.config['DUMP_CACHE_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']), 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_dump_cache()
    cached_images.sort(reverse=True)

    valid_images = [img for img in cached_images if is_image_valid(os.path.join(app.config['DUMP_FOLDER'], img))]
    return render_template('dump.html', images=valid_images)

def emit_gallery_update():
    socketio.emit('update_gallery')

def emit_dump_update():
    socketio.emit('update_dump')

class Watcher(FileSystemEventHandler):
    def on_created(self, event):
        if not event.is_directory and event.src_path.startswith(app.config['UPLOAD_FOLDER']):
            print(f"File created: {event.src_path}")
            update_cache()
            emit_gallery_update()

        if not event.is_directory and event.src_path.startswith(app.config['DUMP_FOLDER']):
            print(f"Dump file created: {event.src_path}")
            update_dump_cache()
            emit_dump_update()

    def on_deleted(self, event):
        if not event.is_directory and event.src_path.startswith(app.config['UPLOAD_FOLDER']):
            print(f"File deleted: {event.src_path}")
            update_cache()
            emit_gallery_update()

        if not event.is_directory and event.src_path.startswith(app.config['DUMP_FOLDER']):
            print(f"Dump file deleted: {event.src_path}")
            update_dump_cache()
            emit_dump_update()

@socketio.on('connect')
def handle_connect():
    if not session.get('logged_in'):
        disconnect()

if __name__ == '__main__':
    event_handler = Watcher()
    observer = Observer()
    observer.schedule(event_handler, path=app.config['UPLOAD_FOLDER'], recursive=False)
    observer.schedule(event_handler, path=app.config['DUMP_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()