Expanded UI

app.py

@@ -9,13 +9,17 @@ 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/Portrait"
+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)
@@ -31,18 +35,22 @@ def init_db():
     )
     ''')
 
-    # Create image_metadata table
+    # (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 TEXT NOT NULL,
+        resolution_x INTEGER NOT NULL,
+        resolution_y INTEGER NOT NULL,
         name TEXT NOT NULL,
         orientation TEXT NOT NULL,
-        discovered_at INTEGER NOT NULL
+        creation_date INTEGER NOT NULL,
+        prompt_data TEXT
     )
     ''')
     
+
+    
     conn.commit()
     conn.close()
     print(f"Database initialized at {DB_PATH}")
@@ -70,7 +78,12 @@ def get_selections():
         cursor = conn.cursor()
         
         cursor.execute('''
-        SELECT * FROM image_selections ORDER BY timestamp DESC
+        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()
@@ -83,6 +96,18 @@ def get_selections():
             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:
@@ -152,8 +177,13 @@ def sync_image_database():
         try:
             with Image.open(full_path) as img:
                 width, height = img.size
-                orientation = 'landscape' if width >= height else 'portrait'
-                images_to_add.append((image_path, res, img_name, orientation, int(time.time())))
+                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
@@ -163,8 +193,8 @@ def sync_image_database():
 
     if images_to_add:
         cursor.executemany('''
-        INSERT INTO image_metadata (path, resolution, name, orientation, discovered_at)
-        VALUES (?, ?, ?, ?, ?)
+        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.")
@@ -255,6 +285,10 @@ class ImageSwipeHandler(BaseHTTPRequestHandler):
             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:
@@ -270,12 +304,16 @@ class ImageSwipeHandler(BaseHTTPRequestHandler):
         parsed_url = urllib.parse.urlparse(self.path)
         path = parsed_url.path
 
-        if path == '/selection':
+        # 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)
@@ -287,6 +325,7 @@ class ImageSwipeHandler(BaseHTTPRequestHandler):
                 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):
@@ -353,7 +392,7 @@ class ImageSwipeHandler(BaseHTTPRequestHandler):
 
             # Base query to get unactioned images
             query = """
-                SELECT meta.path, meta.resolution, meta.name, meta.orientation
+                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
@@ -382,28 +421,27 @@ class ImageSwipeHandler(BaseHTTPRequestHandler):
             # Choose one random image from the filtered list
             chosen_image_row = random.choice(possible_images)
             image_path = chosen_image_row[0]
-            resolution = chosen_image_row[1]
-            image_name = chosen_image_row[2]
-            orientation = chosen_image_row[3]
+            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}")
 
-            # Get file metadata
-            try:
-                file_stat = os.stat(full_image_path)
-                creation_time = file_stat.st_mtime
-                creation_date = datetime.datetime.fromtimestamp(creation_time).strftime('%Y-%m-%d %H:%M:%S')
-            except Exception as e:
-                print(f"DEBUG ERROR getting file creation time: {str(e)}")
-                creation_date = "Unknown"
+
             
             # Return the image path as JSON
             response = {
                 'path': f"/images/{image_path}",
-                'resolution': resolution,
+                'resolution_x': resolution_x,
+                'resolution_y': resolution_y,
+                'resolution': f"{resolution_x}x{resolution_y}",
                 'filename': image_name,
-                'creation_date': creation_date,
+                'creation_date': datetime.datetime.fromtimestamp(creation_ts).strftime('%Y-%m-%d %H:%M:%S'),
+                'prompt_data': prompt_data,
                 'orientation': orientation
             }
             
@@ -473,7 +511,9 @@ class ImageSwipeHandler(BaseHTTPRequestHandler):
         parsed_path = urllib.parse.urlparse(self.path)
         path = parsed_path.path
         
-        if path == "/record-selection":
+        if path == "/selection":
+            self.handle_selection()
+        elif path == "/record-selection":
             self.handle_record_selection()
         elif path == "/update-selection":
             self.handle_update_selection()
@@ -488,6 +528,29 @@ class ImageSwipeHandler(BaseHTTPRequestHandler):
         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
@@ -503,12 +566,12 @@ class ImageSwipeHandler(BaseHTTPRequestHandler):
             action = data.get('action', '')
             
             # Validate the data
-            if not image_path or not resolution or not action:
+            if not image_path or not action:
                 self.send_error(400, "Missing required fields")
                 return
             
-            # Add the selection to the database
-            add_selection(image_path, resolution, action)
+            # Store only image_path & action for compatibility
+            add_selection(image_path, action)
             
             # Return success response
             response = {