Files
swiper/sort_images.py
2025-06-21 21:44:52 +01:00

86 lines
3.3 KiB
Python

import os
import sqlite3
import shutil
# --- CONFIGURATION ---
# The root directory where the original images are stored.
# IMPORTANT: This must match the IMAGE_DIR in your app.py
IMAGE_DIR = "/mnt/secret-items/sd-outputs/Sorted/Images/Portrait"
# The directory where the sorted images will be moved.
OUTPUT_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "swiper-output")
# The path to the selections database.
DB_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "image_selections.db")
def sort_swiped_images():
"""Reads the database, moves swiped files into action-based folders, and cleans the DB."""
print(f"Connecting to database at {DB_PATH}...")
if not os.path.exists(DB_PATH):
print(f"Error: Database file not found at {DB_PATH}")
return
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
try:
# Get all selections from the database
cursor.execute("SELECT id, image_path, action FROM image_selections")
selections = cursor.fetchall()
except sqlite3.OperationalError as e:
print(f"Error querying database: {e}")
print("This might mean the 'image_selections' table doesn't exist or is empty.")
conn.close()
return
if not selections:
print("No swiped images found in the database to sort.")
conn.close()
return
print(f"Found {len(selections)} swiped images to sort.")
moved_count = 0
for selection_id, image_path, action in selections:
# The image_path from the DB might have a leading '/images/' which we need to remove.
# It's safer to handle both cases.
clean_image_path = image_path.replace('/images/', '', 1) if image_path.startswith('/images/') else image_path
source_path = os.path.join(IMAGE_DIR, clean_image_path)
if not os.path.exists(source_path):
print(f"Warning: Source file not found, skipping: {source_path}")
# We should still remove this dangling DB record
cursor.execute("DELETE FROM image_selections WHERE id = ?", (selection_id,))
cursor.execute("DELETE FROM image_metadata WHERE path = ?", (clean_image_path,))
print(f"Removed dangling database entry for {clean_image_path}")
continue
# Create the destination folder based on the action
destination_folder = os.path.join(OUTPUT_DIR, action)
os.makedirs(destination_folder, exist_ok=True)
destination_path = os.path.join(destination_folder, os.path.basename(clean_image_path))
try:
print(f"Moving '{clean_image_path}' to '{action}' folder...")
shutil.move(source_path, destination_path)
moved_count += 1
# If move was successful, remove the records from the database
cursor.execute("DELETE FROM image_selections WHERE id = ?", (selection_id,))
cursor.execute("DELETE FROM image_metadata WHERE path = ?", (clean_image_path,))
print(f"Successfully moved and removed DB entries for {clean_image_path}")
except Exception as e:
print(f"Error moving file {source_path}: {e}")
# Commit all changes to the database
conn.commit()
conn.close()
print(f"\nSorting complete. Moved {moved_count} images.")
if __name__ == "__main__":
sort_swiped_images()