from http.server import HTTPServer, BaseHTTPRequestHandler import os import json import random import mimetypes import urllib.parse import sqlite3 import time import datetime import zipfile import io from PIL import Image # Path to the image directory IMAGE_DIR = "/mnt/secret-items/sd-outputs/Sorted/Images" # Database file path DB_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "image_selections.db") # NOTE: We no longer delete the database on each run. # If schema changes are needed, run a one-time migration script instead. # Initialize the database def init_db(): conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() # Create image_selections table cursor.execute(''' CREATE TABLE IF NOT EXISTS image_selections ( id INTEGER PRIMARY KEY AUTOINCREMENT, image_path TEXT NOT NULL UNIQUE, action TEXT NOT NULL, timestamp INTEGER NOT NULL ) ''') # (Re)create image_metadata table with new schema cursor.execute(''' CREATE TABLE IF NOT EXISTS image_metadata ( id INTEGER PRIMARY KEY AUTOINCREMENT, path TEXT NOT NULL UNIQUE, resolution_x INTEGER NOT NULL, resolution_y INTEGER NOT NULL, name TEXT NOT NULL, orientation TEXT NOT NULL, creation_date INTEGER NOT NULL, prompt_data TEXT ) ''') conn.commit() conn.close() print(f"Database initialized at {DB_PATH}") # Add a selection to the database def add_selection(image_path, action): conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() # Use REPLACE INTO to handle potential duplicates gracefully cursor.execute(''' REPLACE INTO image_selections (image_path, action, timestamp) VALUES (?, ?, ?) ''', (image_path, action, int(time.time()))) conn.commit() conn.close() # Get all selections from the database def get_selections(): print("DEBUG: get_selections() called") try: conn = sqlite3.connect(DB_PATH) conn.row_factory = sqlite3.Row # This enables column access by name cursor = conn.cursor() cursor.execute(''' SELECT sel.id, sel.image_path, sel.action, sel.timestamp, meta.resolution_x, meta.resolution_y, meta.orientation, meta.creation_date, meta.prompt_data, meta.name FROM image_selections sel LEFT JOIN image_metadata meta ON sel.image_path = meta.path ORDER BY sel.timestamp DESC ''') rows = cursor.fetchall() print(f"DEBUG: Fetched {len(rows)} rows from database") # Properly convert SQLite Row objects to dictionaries results = [] for row in rows: item = {} for key in row.keys(): item[key] = row[key] # Ensure resolution exists if 'resolution' not in item or not item['resolution']: # derive resolution from path e.g. 2048x2048 try: path_part = item['image_path'] if path_part.startswith('/images/'): path_part = path_part[8:] res = path_part.split('/')[0] item['resolution'] = res except Exception: item['resolution'] = "unknown" # Ensure orientation exists if 'orientation' not in item or not item['orientation']: try: # Try to determine orientation if not in database image_path = item['image_path'] if image_path.startswith('/images/'): image_path = image_path[8:] full_path = os.path.join(IMAGE_DIR, image_path) with Image.open(full_path) as img: width, height = img.size item['orientation'] = "portrait" if height > width else "landscape" if width > height else "square" except Exception as e: print(f"DEBUG ERROR determining missing orientation: {str(e)}") item['orientation'] = "unknown" results.append(item) print(f"DEBUG: Converted {len(results)} rows to dictionaries") print(f"DEBUG: First result (if any): {results[0] if results else 'None'}") conn.close() return results except Exception as e: print(f"DEBUG ERROR in get_selections(): {str(e)}") # Return empty list on error to prevent client from hanging return [] # Get a list of all image paths that have already been actioned def sync_image_database(): """Scans the image directory and adds any new images to the metadata table.""" print("Syncing image database...") from PIL import Image conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() # Get all image paths already in the database cursor.execute("SELECT path FROM image_metadata") db_images = {row[0] for row in cursor.fetchall()} print(f"Found {len(db_images)} images in the database.") # Find all images on the filesystem disk_images = set() resolutions = [d for d in os.listdir(IMAGE_DIR) if os.path.isdir(os.path.join(IMAGE_DIR, d))] for res in resolutions: res_dir = os.path.join(IMAGE_DIR, res) for img_name in os.listdir(res_dir): if img_name.lower().endswith(('.png', '.jpg', '.jpeg')): disk_images.add(f"{res}/{img_name}") print(f"Found {len(disk_images)} images on disk.") # Determine which images are new new_images = disk_images - db_images print(f"Found {len(new_images)} new images to add to the database.") if not new_images: print("Database is already up-to-date.") conn.close() return # Process and add new images to the database images_to_add = [] total_new_images = len(new_images) processed_count = 0 for image_path in new_images: res, img_name = image_path.split('/', 1) full_path = os.path.join(IMAGE_DIR, image_path) try: with Image.open(full_path) as img: width, height = img.size orientation = 'square' if width == height else ('landscape' if width > height else 'portrait') # Attempt to read prompt info from PNG metadata (PNG only) prompt_text = None if img.format == 'PNG': prompt_text = img.info.get('parameters') or img.info.get('Parameters') creation_ts = int(os.path.getmtime(full_path)) images_to_add.append((image_path, width, height, img_name, orientation, creation_ts, prompt_text)) processed_count += 1 if processed_count % 100 == 0 or processed_count == total_new_images: percentage = (processed_count / total_new_images) * 100 print(f"Processed {processed_count} of {total_new_images} images ({percentage:.2f}%)...", flush=True) except Exception as e: print(f"Could not process image {full_path}: {e}") if images_to_add: cursor.executemany(''' INSERT INTO image_metadata (path, resolution_x, resolution_y, name, orientation, creation_date, prompt_data) VALUES (?, ?, ?, ?, ?, ?, ?) ''', images_to_add) conn.commit() print(f"Successfully added {len(images_to_add)} new images to the database.") conn.close() # Update a selection in the database def update_selection(selection_id, action): conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() # Update the selection cursor.execute(''' UPDATE image_selections SET action = ?, timestamp = ? WHERE id = ? ''', (action, int(time.time()), selection_id)) # Check if a row was affected rows_affected = cursor.rowcount conn.commit() conn.close() return rows_affected > 0 # Delete a selection from the database def delete_selection(selection_id): conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() # Delete the selection cursor.execute(''' DELETE FROM image_selections WHERE id = ? ''', (selection_id,)) # Check if a row was affected rows_affected = cursor.rowcount conn.commit() conn.close() return rows_affected > 0 # Reset the database by deleting all selections def reset_database(): print("DEBUG: Resetting database - deleting all selections") conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() # Delete all selections cursor.execute(''' DELETE FROM image_selections ''') # Get the number of rows affected rows_affected = cursor.rowcount conn.commit() conn.close() print(f"DEBUG: Reset database - deleted {rows_affected} selections") return rows_affected class ImageSwipeHandler(BaseHTTPRequestHandler): # Set response headers for CORS def _set_cors_headers(self): self.send_header('Access-Control-Allow-Origin', '*') self.send_header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS') self.send_header('Access-Control-Allow-Headers', 'Content-Type') def do_GET(self): # Parse the URL parsed_url = urllib.parse.urlparse(self.path) path = parsed_url.path # Handle different paths if path == '/': self.serve_file('index.html', 'text/html') elif path == '/history': self.serve_file('history.html', 'text/html') elif path == '/styles.css': self.serve_file('styles.css', 'text/css') elif path == '/script.js': self.serve_file('script.js', 'application/javascript') elif path == '/random-image': self.serve_random_image() elif path == '/selections': self.serve_selections() elif path.startswith('/images/'): self.serve_image(path[8:]) elif path == '/favicon.ico': # Silently ignore favicon requests self.send_response(204) self.end_headers() elif path.startswith('/download-selected'): self.handle_download_selected() else: # Try to serve as a static file if path.startswith('/'): path = path[1:] # Remove leading slash try: self.serve_file(path) except: self.send_error(404, "File not found") def do_POST(self): parsed_url = urllib.parse.urlparse(self.path) path = parsed_url.path # Debug: log every POST path print(f"DEBUG: do_POST received path='{path}'") # Accept /selection paths if path.startswith('/selection'): try: content_length = int(self.headers['Content-Length']) post_data = self.rfile.read(content_length) data = json.loads(post_data) print(f"DEBUG: Received selection POST: {data}") add_selection(data['image_path'], data['action']) self.send_response(200) self.send_header('Content-type', 'application/json') self._set_cors_headers() self.end_headers() self.wfile.write(json.dumps({'status': 'success'}).encode()) except Exception as e: print(f"ERROR in do_POST /selection: {e}") self.send_error(500, f"Server error processing selection: {e}") else: print(f"DEBUG: Unknown POST path '{path}'") self.send_error(404, "Endpoint not found") def do_OPTIONS(self): self.send_response(204) self._set_cors_headers() self.end_headers() def serve_file(self, file_path, content_type=None): try: with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), file_path), 'rb') as file: content = file.read() self.send_response(200) # Set the content type based on parameter or guess from file extension if not content_type: content_type, _ = mimetypes.guess_type(file_path) if content_type: self.send_header('Content-type', content_type) else: self.send_header('Content-type', 'application/octet-stream') self._set_cors_headers() self.send_header('Content-length', len(content)) self.end_headers() self.wfile.write(content) except FileNotFoundError: self.send_error(404, f"File not found: {file_path}") def serve_image(self, image_path): try: # Decode URL-encoded path image_path = urllib.parse.unquote(image_path) full_path = os.path.join(IMAGE_DIR, image_path) with open(full_path, 'rb') as file: content = file.read() self.send_response(200) # Set the content type based on file extension content_type, _ = mimetypes.guess_type(full_path) if content_type: self.send_header('Content-type', content_type) else: self.send_header('Content-type', 'application/octet-stream') self._set_cors_headers() self.send_header('Content-length', len(content)) self.end_headers() self.wfile.write(content) except FileNotFoundError: self.send_error(404, f"Image not found: {image_path}") def serve_random_image(self): try: parsed_url = urllib.parse.urlparse(self.path) query_params = urllib.parse.parse_qs(parsed_url.query) orientation_filter = query_params.get('orientation', ['all'])[0] conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() # Base query to get unactioned images query = """ SELECT meta.path, meta.resolution_x, meta.resolution_y, meta.name, meta.orientation, meta.creation_date, meta.prompt_data FROM image_metadata meta LEFT JOIN image_selections sel ON meta.path = sel.image_path WHERE sel.image_path IS NULL """ # Add orientation filter if specified params = () if orientation_filter != 'all': query += " AND meta.orientation = ?" params = (orientation_filter,) cursor.execute(query, params) possible_images = cursor.fetchall() conn.close() if not possible_images: print("DEBUG: No matching unactioned images found.") response = {'message': 'No more images available for this filter.'} self.send_response(200) self.send_header('Content-type', 'application/json') self._set_cors_headers() self.end_headers() self.wfile.write(json.dumps(response).encode()) return # Choose one random image from the filtered list chosen_image_row = random.choice(possible_images) image_path = chosen_image_row[0] resolution_x = chosen_image_row[1] resolution_y = chosen_image_row[2] image_name = chosen_image_row[3] orientation = chosen_image_row[4] creation_ts = chosen_image_row[5] prompt_data = chosen_image_row[6] full_image_path = os.path.join(IMAGE_DIR, image_path) print(f"DEBUG: Serving image: {image_path}") # Return the image path as JSON response = { 'path': f"/images/{image_path}", 'resolution_x': resolution_x, 'resolution_y': resolution_y, 'resolution': f"{resolution_x}x{resolution_y}", 'filename': image_name, 'creation_date': datetime.datetime.fromtimestamp(creation_ts).strftime('%Y-%m-%d %H:%M:%S'), 'prompt_data': prompt_data, 'orientation': orientation } self.send_response(200) self.send_header('Content-type', 'application/json') self._set_cors_headers() self.end_headers() self.wfile.write(json.dumps(response).encode()) except Exception as e: print(f"FATAL ERROR in serve_random_image: {e}") self.send_error(500, f"Error serving random image: {str(e)}") def serve_resolutions(self): try: # Get all resolution directories resolutions = [d for d in os.listdir(IMAGE_DIR) if os.path.isdir(os.path.join(IMAGE_DIR, d))] # Return the resolutions as JSON response = { 'resolutions': resolutions } self.send_response(200) self.send_header('Content-type', 'application/json') self._set_cors_headers() self.end_headers() self.wfile.write(json.dumps(response).encode()) except Exception as e: self.send_error(500, f"Error serving resolutions: {str(e)}") def serve_selections(self): print("DEBUG: serve_selections() called") try: # Get all selections from the database selections = get_selections() # Return the selections as JSON response = { 'selections': selections } # Debug the response before sending print(f"DEBUG: Response has {len(selections)} selections") # Try to serialize to JSON to catch any serialization errors try: response_json = json.dumps(response) print(f"DEBUG: JSON serialization successful, length: {len(response_json)}") except Exception as json_err: print(f"DEBUG ERROR in JSON serialization: {str(json_err)}") # If there's an error in serialization, send a simpler response response = {'selections': [], 'error': 'JSON serialization error'} response_json = json.dumps(response) self.send_response(200) self.send_header('Content-type', 'application/json') self._set_cors_headers() self.end_headers() self.wfile.write(response_json.encode()) print("DEBUG: Response sent successfully") except Exception as e: print(f"DEBUG ERROR in serve_selections(): {str(e)}") self.send_error(500, f"Error serving selections: {str(e)}") def do_POST(self): # Parse the URL path parsed_path = urllib.parse.urlparse(self.path) path = parsed_path.path if path == "/selection": self.handle_selection() elif path == "/record-selection": self.handle_record_selection() elif path == "/update-selection": self.handle_update_selection() elif path == "/delete-selection": self.handle_delete_selection() elif self.path == '/reset-database': self.handle_reset_database() return elif self.path.startswith('/download-selected'): self.handle_download_selected() return else: self.send_error(404, "Not found") def handle_selection(self): """Handle legacy /selection POST with image_path and action""" try: content_length = int(self.headers['Content-Length']) post_data = self.rfile.read(content_length) data = json.loads(post_data) image_path = data.get('image_path') action = data.get('action') if not image_path or not action: self.send_error(400, "Missing required fields") return add_selection(image_path, action) self.send_response(200) self.send_header('Content-type', 'application/json') self._set_cors_headers() self.end_headers() self.wfile.write(json.dumps({'status': 'success'}).encode()) except Exception as e: print(f"ERROR in handle_selection: {e}") self.send_error(500, f"Server error: {e}") def handle_record_selection(self): try: # Get the content length content_length = int(self.headers['Content-Length']) # Read the request body post_data = self.rfile.read(content_length).decode('utf-8') data = json.loads(post_data) # Extract the required fields image_path = data.get('path', '').replace('/images/', '') resolution = data.get('resolution', '') action = data.get('action', '') # Validate the data if not image_path or not action: self.send_error(400, "Missing required fields") return # Store only image_path & action for compatibility add_selection(image_path, action) # Return success response response = { 'success': True, 'message': f"Selection recorded: {action} for {image_path}" } self.send_response(200) self.send_header('Content-type', 'application/json') self._set_cors_headers() self.end_headers() self.wfile.write(json.dumps(response).encode()) except Exception as e: self.send_error(500, f"Error recording selection: {str(e)}") def handle_update_selection(self): try: # Get the content length content_length = int(self.headers['Content-Length']) # Read the request body post_data = self.rfile.read(content_length).decode('utf-8') data = json.loads(post_data) # Extract the required fields selection_id = data.get('id') action = data.get('action', '') # Validate the data if not selection_id or not action: self.send_error(400, "Missing required fields") return # Update the selection in the database success = update_selection(selection_id, action) if not success: self.send_error(404, f"Selection with ID {selection_id} not found") return # Return success response response = { 'success': True, 'message': f"Selection updated: ID {selection_id} to {action}" } self.send_response(200) self.send_header('Content-type', 'application/json') self._set_cors_headers() self.end_headers() self.wfile.write(json.dumps(response).encode()) except Exception as e: self.send_error(500, f"Error updating selection: {str(e)}") def handle_delete_selection(self): try: # Get the content length content_length = int(self.headers['Content-Length']) # Read the request body post_data = self.rfile.read(content_length).decode('utf-8') data = json.loads(post_data) # Extract the required fields selection_id = data.get('id') # Validate the data if not selection_id: self.send_error(400, "Missing selection ID") return # Delete the selection from the database success = delete_selection(selection_id) if not success: self.send_error(404, f"Selection with ID {selection_id} not found") return # Return success response response = { 'success': True, 'message': f"Selection deleted: ID {selection_id}" } self.send_response(200) self.send_header('Content-type', 'application/json') self._set_cors_headers() self.end_headers() self.wfile.write(json.dumps(response).encode()) except Exception as e: print(f"DEBUG ERROR in handle_delete_selection(): {str(e)}") self.send_error(500, f"Error deleting selection: {str(e)}") def handle_reset_database(self): try: reset_database() self.send_response(200) self.send_header('Content-type', 'application/json') self.send_header('Access-Control-Allow-Origin', '*') self.end_headers() self.wfile.write(json.dumps({'success': True, 'message': 'Database reset successfully'}).encode()) except Exception as e: print(f"DEBUG ERROR in handle_reset_database: {str(e)}") self.send_response(500) self.send_header('Content-type', 'application/json') self.send_header('Access-Control-Allow-Origin', '*') self.end_headers() self.wfile.write(json.dumps({'success': False, 'message': f'Error: {str(e)}'}).encode()) def handle_download_selected(self): try: # Parse the query parameters to get the selected image paths query_components = urllib.parse.parse_qs(urllib.parse.urlparse(self.path).query) image_paths = query_components.get('paths', []) if not image_paths: self.send_response(400) self.send_header('Content-type', 'application/json') self.send_header('Access-Control-Allow-Origin', '*') self.end_headers() self.wfile.write(json.dumps({'success': False, 'message': 'No image paths provided'}).encode()) return # Create a zip file in memory zip_buffer = io.BytesIO() with zipfile.ZipFile(zip_buffer, 'a', zipfile.ZIP_DEFLATED, False) as zip_file: for path in image_paths: # Remove the /images/ prefix if path.startswith('/images/'): path = path[8:] full_path = os.path.join(IMAGE_DIR, path) if os.path.exists(full_path): # Add the file to the zip with just the filename (no directory structure) filename = os.path.basename(path) zip_file.write(full_path, filename) # Seek to the beginning of the buffer zip_buffer.seek(0) # Send the zip file as a response self.send_response(200) self.send_header('Content-type', 'application/zip') self.send_header('Content-Disposition', 'attachment; filename="selected_images.zip"') self.send_header('Access-Control-Allow-Origin', '*') self.end_headers() self.wfile.write(zip_buffer.getvalue()) except Exception as e: print(f"DEBUG ERROR in handle_download_selected: {str(e)}") self.send_response(500) self.send_header('Content-type', 'application/json') self.send_header('Access-Control-Allow-Origin', '*') self.end_headers() self.wfile.write(json.dumps({'success': False, 'message': f'Error: {str(e)}'}).encode()) def do_OPTIONS(self): # Handle preflight requests for CORS self.send_response(200) self._set_cors_headers() self.end_headers() def run(server_class=HTTPServer, handler_class=ImageSwipeHandler, port=8000): # Initialize the database init_db() # Ensure the 'images' directory exists if not os.path.exists(IMAGE_DIR): os.makedirs(IMAGE_DIR) # Sync the image database on startup sync_image_database() server_address = ('', port) httpd = server_class(server_address, handler_class) print(f"Starting server on port {port}...") print(f"Image directory: {IMAGE_DIR}") print(f"Database: {DB_PATH}") httpd.serve_forever() if __name__ == "__main__": run()