Compare commits

12 Commits

Author SHA1 Message Date
0c2504ff83 Fix caption cache timing, commit path tracking, and clean up start script
- Move caption cache refresh before UI render so indicators show on load
- Return actual dest paths from commit_batch/commit_global to fix
  caption-on-apply silently failing when files are renamed on collision
- Simplify start.sh to only run NiceGUI (remove Streamlit)
- Add requests to requirements.txt

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-02-08 18:59:05 +01:00
43772aba68 caption fix 2026-01-28 16:15:42 +01:00
145368692e caption 2026-01-28 16:06:56 +01:00
bf1134e47f claude 2026-01-28 15:42:14 +01:00
7580036c9d Update engine.py 2026-01-23 13:52:39 +01:00
47a75b428e Update gallery_app.py 2026-01-23 13:52:25 +01:00
b91a2f0a31 Update engine.py 2026-01-23 13:42:43 +01:00
66795471a8 Update gallery_app.py 2026-01-23 13:42:28 +01:00
67acb8e08a Update gallery_app.py 2026-01-23 13:35:55 +01:00
af0cc52d89 Update gallery_app.py 2026-01-23 13:27:04 +01:00
3669814731 Add engine.py 2026-01-23 13:26:45 +01:00
a8edc251a2 Delete engine.py 2026-01-23 13:26:23 +01:00
7 changed files with 1749 additions and 271 deletions

61
CLAUDE.md Normal file
View File

@@ -0,0 +1,61 @@
# CLAUDE.md
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
## Project Overview
Turbo Sorter Pro v12.5 - A dual-interface image organization tool combining Streamlit (admin dashboard) and NiceGUI (gallery interface) for managing large image collections through time-sync matching, ID collision resolution, category-based sorting, and gallery tagging with pairing capabilities.
## Commands
```bash
# Install dependencies
pip install -r requirements.txt
# Run Streamlit dashboard (port 8501)
streamlit run app.py --server.port=8501 --server.address=0.0.0.0
# Run NiceGUI gallery (port 8080)
python3 gallery_app.py
# Both services (container startup)
./start.sh
# Syntax check all Python files
python3 -m py_compile *.py
```
## Architecture
### Dual-Framework Design
- **Streamlit (app.py, port 8501)**: Administrative dashboard with 5 modular tabs for management workflows
- **NiceGUI (gallery_app.py, port 8080)**: Modern gallery interface for image tagging and pairing operations
- **Shared Backend**: Both UIs use `SorterEngine` (engine.py) and the same SQLite database
### Core Components
| File | Purpose |
|------|---------|
| `engine.py` | Static `SorterEngine` class - all DB operations, file handling, image compression |
| `gallery_app.py` | NiceGUI gallery with `AppState` class for centralized state management |
| `app.py` | Streamlit entry point, loads tab modules |
| `tab_*.py` | Independent tab modules for each workflow |
### Database
SQLite at `/app/sorter_database.db` with tables: profiles, folder_ids, categories, staging_area, processed_log, folder_tags, profile_categories, pairing_settings.
### Tab Workflows
1. **Time-Sync Discovery** - Match images by timestamp across folders
2. **ID Review** - Resolve ID collisions between target/control folders
3. **Unused Archive** - Manage rejected image pairs
4. **Category Sorter** - One-to-many categorization
5. **Gallery Staged** - Grid-based tagging with Gallery/Pairing dual modes
## Key Patterns
- **ID Format**: `id001_`, `id002_` (zero-padded 3-digit prefix)
- **Staging Pattern**: Two-phase commit (stage → commit) with undo support
- **Image Formats**: .jpg, .jpeg, .png, .webp, .bmp, .tiff
- **Compression**: WebP with ThreadPoolExecutor (8 workers)
- **Permissions**: chmod 0o777 applied to committed files
- **Default Paths**: `/storage` when not configured

Binary file not shown.

Binary file not shown.

674
engine.py
View File

@@ -1,12 +1,31 @@
import os import os
import shutil import shutil
import sqlite3 import sqlite3
import base64
import requests
from datetime import datetime
from contextlib import contextmanager
from PIL import Image from PIL import Image
from io import BytesIO from io import BytesIO
class SorterEngine: class SorterEngine:
DB_PATH = "/app/sorter_database.db" DB_PATH = "/app/sorter_database.db"
@staticmethod
@contextmanager
def get_db():
"""Context manager for database connections.
Ensures proper commit/rollback and always closes connection."""
conn = sqlite3.connect(SorterEngine.DB_PATH)
try:
yield conn
conn.commit()
except Exception:
conn.rollback()
raise
finally:
conn.close()
# --- 1. DATABASE INITIALIZATION --- # --- 1. DATABASE INITIALIZATION ---
@staticmethod @staticmethod
def init_db(): def init_db():
@@ -28,29 +47,84 @@ class SorterEngine:
(source_path TEXT PRIMARY KEY, category TEXT, action_type TEXT)''') (source_path TEXT PRIMARY KEY, category TEXT, action_type TEXT)''')
# --- NEW: FOLDER TAGS TABLE (persists tags by folder) --- # --- NEW: FOLDER TAGS TABLE (persists tags by folder) ---
cursor.execute('''CREATE TABLE IF NOT EXISTS folder_tags # Check if old schema exists (without profile column) and migrate
(folder_path TEXT, filename TEXT, category TEXT, tag_index INTEGER, cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='folder_tags'")
PRIMARY KEY (folder_path, filename))''') if cursor.fetchone():
cursor.execute("PRAGMA table_info(folder_tags)")
columns = [row[1] for row in cursor.fetchall()]
if 'profile' not in columns:
# Migrate: drop old table and recreate with profile column
cursor.execute("DROP TABLE folder_tags")
conn.commit()
cursor.execute('''CREATE TABLE IF NOT EXISTS folder_tags
(profile TEXT, folder_path TEXT, filename TEXT, category TEXT, tag_index INTEGER,
PRIMARY KEY (profile, folder_path, filename))''')
# --- NEW: PROFILE CATEGORIES TABLE (each profile has its own categories) --- # --- NEW: PROFILE CATEGORIES TABLE (each profile has its own categories) ---
cursor.execute('''CREATE TABLE IF NOT EXISTS profile_categories cursor.execute('''CREATE TABLE IF NOT EXISTS profile_categories
(profile TEXT, category TEXT, PRIMARY KEY (profile, category))''') (profile TEXT, category TEXT, PRIMARY KEY (profile, category))''')
# --- NEW: PAIRING SETTINGS TABLE ---
cursor.execute('''CREATE TABLE IF NOT EXISTS pairing_settings
(profile TEXT PRIMARY KEY,
adjacent_folder TEXT,
main_category TEXT,
adj_category TEXT,
main_output TEXT,
adj_output TEXT,
time_window INTEGER)''')
# Seed categories if empty (legacy table) # Seed categories if empty (legacy table)
cursor.execute("SELECT COUNT(*) FROM categories") cursor.execute("SELECT COUNT(*) FROM categories")
if cursor.fetchone()[0] == 0: if cursor.fetchone()[0] == 0:
for cat in ["_TRASH", "control", "Default", "Action", "Solo"]: for cat in ["_TRASH", "control", "Default", "Action", "Solo"]:
cursor.execute("INSERT OR IGNORE INTO categories VALUES (?)", (cat,)) cursor.execute("INSERT OR IGNORE INTO categories VALUES (?)", (cat,))
# --- CAPTION TABLES ---
# Per-category prompt templates
cursor.execute('''CREATE TABLE IF NOT EXISTS category_prompts
(profile TEXT, category TEXT, prompt_template TEXT,
PRIMARY KEY (profile, category))''')
# Stored captions
cursor.execute('''CREATE TABLE IF NOT EXISTS image_captions
(image_path TEXT PRIMARY KEY, caption TEXT, model TEXT,
generated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
# Caption API settings per profile
cursor.execute('''CREATE TABLE IF NOT EXISTS caption_settings
(profile TEXT PRIMARY KEY,
api_endpoint TEXT DEFAULT 'http://localhost:8080/v1/chat/completions',
model_name TEXT DEFAULT 'local-model',
max_tokens INTEGER DEFAULT 300,
temperature REAL DEFAULT 0.7,
timeout_seconds INTEGER DEFAULT 60,
batch_size INTEGER DEFAULT 4)''')
# --- PERFORMANCE INDEXES ---
# Index for staging_area queries filtered by category
cursor.execute("CREATE INDEX IF NOT EXISTS idx_staging_category ON staging_area(target_category)")
# Index for folder_tags queries filtered by profile and folder_path
cursor.execute("CREATE INDEX IF NOT EXISTS idx_folder_tags_profile ON folder_tags(profile, folder_path)")
# Index for profile_categories lookups
cursor.execute("CREATE INDEX IF NOT EXISTS idx_profile_categories ON profile_categories(profile)")
# Index for caption lookups by image path
cursor.execute("CREATE INDEX IF NOT EXISTS idx_image_captions ON image_captions(image_path)")
conn.commit() conn.commit()
conn.close() conn.close()
# --- 2. PROFILE & PATH MANAGEMENT --- # --- 2. PROFILE & PATH MANAGEMENT ---
@staticmethod @staticmethod
def save_tab_paths(profile_name, t1_t=None, t2_t=None, t2_c=None, t4_s=None, t4_o=None, mode=None, t5_s=None, t5_o=None): def save_tab_paths(profile_name, t1_t=None, t2_t=None, t2_c=None, t4_s=None, t4_o=None, mode=None, t5_s=None, t5_o=None,
pair_adjacent_folder=None, pair_main_category=None, pair_adj_category=None,
pair_main_output=None, pair_adj_output=None, pair_time_window=None):
"""Updates specific tab paths in the database while preserving others.""" """Updates specific tab paths in the database while preserving others."""
conn = sqlite3.connect(SorterEngine.DB_PATH) conn = sqlite3.connect(SorterEngine.DB_PATH)
cursor = conn.cursor() cursor = conn.cursor()
# Save main profile settings
cursor.execute("SELECT * FROM profiles WHERE name = ?", (profile_name,)) cursor.execute("SELECT * FROM profiles WHERE name = ?", (profile_name,))
row = cursor.fetchone() row = cursor.fetchone()
@@ -70,6 +144,38 @@ class SorterEngine:
t5_o if t5_o is not None else row[8] t5_o if t5_o is not None else row[8]
) )
cursor.execute("INSERT OR REPLACE INTO profiles VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", new_values) cursor.execute("INSERT OR REPLACE INTO profiles VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", new_values)
# Save pairing settings if any are provided
if any(x is not None for x in [pair_adjacent_folder, pair_main_category, pair_adj_category,
pair_main_output, pair_adj_output, pair_time_window]):
# Ensure table exists
cursor.execute('''CREATE TABLE IF NOT EXISTS pairing_settings
(profile TEXT PRIMARY KEY,
adjacent_folder TEXT,
main_category TEXT,
adj_category TEXT,
main_output TEXT,
adj_output TEXT,
time_window INTEGER)''')
# Get existing values
cursor.execute("SELECT * FROM pairing_settings WHERE profile = ?", (profile_name,))
pair_row = cursor.fetchone()
if not pair_row:
pair_row = (profile_name, "", "control", "control", "/storage", "/storage", 60)
pair_values = (
profile_name,
pair_adjacent_folder if pair_adjacent_folder is not None else pair_row[1],
pair_main_category if pair_main_category is not None else pair_row[2],
pair_adj_category if pair_adj_category is not None else pair_row[3],
pair_main_output if pair_main_output is not None else pair_row[4],
pair_adj_output if pair_adj_output is not None else pair_row[5],
pair_time_window if pair_time_window is not None else pair_row[6]
)
cursor.execute("INSERT OR REPLACE INTO pairing_settings VALUES (?, ?, ?, ?, ?, ?, ?)", pair_values)
conn.commit() conn.commit()
conn.close() conn.close()
@staticmethod @staticmethod
@@ -100,17 +206,50 @@ class SorterEngine:
@staticmethod @staticmethod
def load_profiles(): def load_profiles():
"""Loads all workspace presets.""" """Loads all workspace presets including pairing settings.
Uses LEFT JOIN to fetch all data in a single query (fixes N+1 problem)."""
conn = sqlite3.connect(SorterEngine.DB_PATH) conn = sqlite3.connect(SorterEngine.DB_PATH)
cursor = conn.cursor() cursor = conn.cursor()
cursor.execute("SELECT * FROM profiles")
# Ensure pairing_settings table exists before JOIN
cursor.execute('''CREATE TABLE IF NOT EXISTS pairing_settings
(profile TEXT PRIMARY KEY,
adjacent_folder TEXT,
main_category TEXT,
adj_category TEXT,
main_output TEXT,
adj_output TEXT,
time_window INTEGER)''')
# Single query with LEFT JOIN - eliminates N+1 queries
cursor.execute('''
SELECT p.name, p.tab1_target, p.tab2_target, p.tab2_control,
p.tab4_source, p.tab4_out, p.mode, p.tab5_source, p.tab5_out,
ps.adjacent_folder, ps.main_category, ps.adj_category,
ps.main_output, ps.adj_output, ps.time_window
FROM profiles p
LEFT JOIN pairing_settings ps ON p.name = ps.profile
''')
rows = cursor.fetchall() rows = cursor.fetchall()
profiles = {}
for r in rows:
profile_name = r[0]
profiles[profile_name] = {
"tab1_target": r[1], "tab2_target": r[2], "tab2_control": r[3],
"tab4_source": r[4], "tab4_out": r[5], "mode": r[6],
"tab5_source": r[7], "tab5_out": r[8],
# Pairing settings from JOIN (with defaults for NULL)
"pair_adjacent_folder": r[9] or "",
"pair_main_category": r[10] or "control",
"pair_adj_category": r[11] or "control",
"pair_main_output": r[12] or "/storage",
"pair_adj_output": r[13] or "/storage",
"pair_time_window": r[14] or 60
}
conn.close() conn.close()
return {r[0]: { return profiles
"tab1_target": r[1], "tab2_target": r[2], "tab2_control": r[3],
"tab4_source": r[4], "tab4_out": r[5], "mode": r[6],
"tab5_source": r[7], "tab5_out": r[8]
} for r in rows}
# --- 3. CATEGORY MANAGEMENT (Profile-based) --- # --- 3. CATEGORY MANAGEMENT (Profile-based) ---
@staticmethod @staticmethod
@@ -186,20 +325,42 @@ class SorterEngine:
# --- 4. IMAGE & ID OPERATIONS --- # --- 4. IMAGE & ID OPERATIONS ---
@staticmethod @staticmethod
def get_images(path, recursive=False): def get_images(path, recursive=False, exclude_paths=None):
"""Image scanner with optional recursive subfolder support.""" """Image scanner with optional recursive subfolder support.
Args:
path: Directory to scan
recursive: Whether to scan subdirectories
exclude_paths: List of paths to exclude from scanning
"""
exts = ('.jpg', '.jpeg', '.png', '.webp', '.bmp', '.tiff') exts = ('.jpg', '.jpeg', '.png', '.webp', '.bmp', '.tiff')
if not path or not os.path.exists(path): return [] if not path or not os.path.exists(path): return []
exclude_paths = exclude_paths or []
# Normalize exclude paths
exclude_normalized = [os.path.normpath(os.path.abspath(p)) for p in exclude_paths]
image_list = [] image_list = []
if recursive: if recursive:
for root, _, files in os.walk(path): for root, dirs, files in os.walk(path):
# Skip the trash folder from scanning # Skip the trash folder from scanning
if "_DELETED" in root: continue if "_DELETED" in root: continue
# Skip excluded paths
root_normalized = os.path.normpath(os.path.abspath(root))
if any(root_normalized.startswith(exc) or exc.startswith(root_normalized) for exc in exclude_normalized):
# Remove excluded dirs from dirs to prevent descending into them
dirs[:] = [d for d in dirs if os.path.normpath(os.path.abspath(os.path.join(root, d))) not in exclude_normalized]
if root_normalized in exclude_normalized:
continue
for f in files: for f in files:
if f.lower().endswith(exts): image_list.append(os.path.join(root, f)) if f.lower().endswith(exts):
image_list.append(os.path.join(root, f))
else: else:
for f in os.listdir(path): for f in os.listdir(path):
if f.lower().endswith(exts): image_list.append(os.path.join(path, f)) if f.lower().endswith(exts):
image_list.append(os.path.join(path, f))
return sorted(image_list) return sorted(image_list)
@staticmethod @staticmethod
@@ -259,60 +420,55 @@ class SorterEngine:
@staticmethod @staticmethod
def stage_image(original_path, category, new_name): def stage_image(original_path, category, new_name):
"""Records a pending rename/move in the database.""" """Records a pending rename/move in the database."""
conn = sqlite3.connect(SorterEngine.DB_PATH) with SorterEngine.get_db() as conn:
cursor = conn.cursor() cursor = conn.cursor()
cursor.execute("INSERT OR REPLACE INTO staging_area VALUES (?, ?, ?, 1)", (original_path, category, new_name)) cursor.execute("INSERT OR REPLACE INTO staging_area VALUES (?, ?, ?, 1)", (original_path, category, new_name))
conn.commit()
conn.close()
@staticmethod @staticmethod
def clear_staged_item(original_path): def clear_staged_item(original_path):
"""Removes an item from the pending staging area.""" """Removes an item from the pending staging area."""
conn = sqlite3.connect(SorterEngine.DB_PATH) with SorterEngine.get_db() as conn:
cursor = conn.cursor() cursor = conn.cursor()
cursor.execute("DELETE FROM staging_area WHERE original_path = ?", (original_path,)) cursor.execute("DELETE FROM staging_area WHERE original_path = ?", (original_path,))
conn.commit()
conn.close()
@staticmethod @staticmethod
def clear_staging_area(): def clear_staging_area():
"""Clears all items from the staging area.""" """Clears all items from the staging area."""
conn = sqlite3.connect(SorterEngine.DB_PATH) with SorterEngine.get_db() as conn:
cursor = conn.cursor() cursor = conn.cursor()
cursor.execute("DELETE FROM staging_area") cursor.execute("DELETE FROM staging_area")
conn.commit()
conn.close()
@staticmethod @staticmethod
def get_staged_data(): def get_staged_data():
"""Retrieves current tagged/staged images.""" """Retrieves current tagged/staged images."""
conn = sqlite3.connect(SorterEngine.DB_PATH) with SorterEngine.get_db() as conn:
cursor = conn.cursor() cursor = conn.cursor()
cursor.execute("SELECT * FROM staging_area") cursor.execute("SELECT * FROM staging_area")
rows = cursor.fetchall() rows = cursor.fetchall()
conn.close() # FIXED: Added "marked": r[3] to the dictionary
# FIXED: Added "marked": r[3] to the dictionary return {r[0]: {"cat": r[1], "name": r[2], "marked": r[3]} for r in rows}
return {r[0]: {"cat": r[1], "name": r[2], "marked": r[3]} for r in rows}
@staticmethod @staticmethod
def commit_global(output_root, cleanup_mode, operation="Copy", source_root=None, profile=None): def commit_global(output_root, cleanup_mode, operation="Copy", source_root=None, profile=None):
"""Commits ALL staged files and fixes permissions.""" """Commits ALL staged files and fixes permissions.
Returns dict mapping original_path -> {dest, cat} for committed files."""
data = SorterEngine.get_staged_data() data = SorterEngine.get_staged_data()
committed = {}
# Save folder tags BEFORE processing (so we can restore them later) # Save folder tags BEFORE processing (so we can restore them later)
if source_root: if source_root:
SorterEngine.save_folder_tags(source_root, profile) SorterEngine.save_folder_tags(source_root, profile)
conn = sqlite3.connect(SorterEngine.DB_PATH) conn = sqlite3.connect(SorterEngine.DB_PATH)
cursor = conn.cursor() cursor = conn.cursor()
if not os.path.exists(output_root): os.makedirs(output_root, exist_ok=True) if not os.path.exists(output_root): os.makedirs(output_root, exist_ok=True)
# 1. Process all Staged Items # 1. Process all Staged Items
for old_p, info in data.items(): for old_p, info in data.items():
if os.path.exists(old_p): if os.path.exists(old_p):
final_dst = os.path.join(output_root, info['name']) final_dst = os.path.join(output_root, info['name'])
if os.path.exists(final_dst): if os.path.exists(final_dst):
root, ext = os.path.splitext(info['name']) root, ext = os.path.splitext(info['name'])
c = 1 c = 1
@@ -324,12 +480,15 @@ class SorterEngine:
shutil.copy2(old_p, final_dst) shutil.copy2(old_p, final_dst)
else: else:
shutil.move(old_p, final_dst) shutil.move(old_p, final_dst)
# --- FIX PERMISSIONS --- # --- FIX PERMISSIONS ---
SorterEngine.fix_permissions(final_dst) SorterEngine.fix_permissions(final_dst)
# Track actual destination
committed[old_p] = {"dest": final_dst, "cat": info['cat']}
# Log History # Log History
cursor.execute("INSERT OR REPLACE INTO processed_log VALUES (?, ?, ?)", cursor.execute("INSERT OR REPLACE INTO processed_log VALUES (?, ?, ?)",
(old_p, info['cat'], operation)) (old_p, info['cat'], operation))
# 2. Global Cleanup # 2. Global Cleanup
@@ -341,16 +500,17 @@ class SorterEngine:
unused_dir = os.path.join(source_root, "unused") unused_dir = os.path.join(source_root, "unused")
os.makedirs(unused_dir, exist_ok=True) os.makedirs(unused_dir, exist_ok=True)
dest_unused = os.path.join(unused_dir, os.path.basename(img_p)) dest_unused = os.path.join(unused_dir, os.path.basename(img_p))
shutil.move(img_p, dest_unused) shutil.move(img_p, dest_unused)
SorterEngine.fix_permissions(dest_unused) SorterEngine.fix_permissions(dest_unused)
elif cleanup_mode == "Delete": elif cleanup_mode == "Delete":
os.remove(img_p) os.remove(img_p)
cursor.execute("DELETE FROM staging_area") cursor.execute("DELETE FROM staging_area")
conn.commit() conn.commit()
conn.close() conn.close()
return committed
# --- 6. CORE UTILITIES (SYNC & UNDO) --- # --- 6. CORE UTILITIES (SYNC & UNDO) ---
@staticmethod @staticmethod
@@ -462,21 +622,23 @@ class SorterEngine:
@staticmethod @staticmethod
def commit_batch(file_list, output_root, cleanup_mode, operation="Copy"): def commit_batch(file_list, output_root, cleanup_mode, operation="Copy"):
"""Commits files and fixes permissions.""" """Commits files and fixes permissions.
Returns dict mapping original_path -> actual_dest_path for committed files."""
data = SorterEngine.get_staged_data() data = SorterEngine.get_staged_data()
conn = sqlite3.connect(SorterEngine.DB_PATH) conn = sqlite3.connect(SorterEngine.DB_PATH)
cursor = conn.cursor() cursor = conn.cursor()
committed = {}
if not os.path.exists(output_root): os.makedirs(output_root, exist_ok=True) if not os.path.exists(output_root): os.makedirs(output_root, exist_ok=True)
for file_path in file_list: for file_path in file_list:
if not os.path.exists(file_path): continue if not os.path.exists(file_path): continue
# --- CASE A: Tagged --- # --- CASE A: Tagged ---
if file_path in data and data[file_path]['marked']: if file_path in data and data[file_path]['marked']:
info = data[file_path] info = data[file_path]
final_dst = os.path.join(output_root, info['name']) final_dst = os.path.join(output_root, info['name'])
# Collision Check # Collision Check
if os.path.exists(final_dst): if os.path.exists(final_dst):
root, ext = os.path.splitext(info['name']) root, ext = os.path.splitext(info['name'])
@@ -484,7 +646,7 @@ class SorterEngine:
while os.path.exists(final_dst): while os.path.exists(final_dst):
final_dst = os.path.join(output_root, f"{root}_{c}{ext}") final_dst = os.path.join(output_root, f"{root}_{c}{ext}")
c += 1 c += 1
# Perform Action # Perform Action
if operation == "Copy": if operation == "Copy":
shutil.copy2(file_path, final_dst) shutil.copy2(file_path, final_dst)
@@ -494,26 +656,30 @@ class SorterEngine:
# --- FIX PERMISSIONS --- # --- FIX PERMISSIONS ---
SorterEngine.fix_permissions(final_dst) SorterEngine.fix_permissions(final_dst)
# Track actual destination
committed[file_path] = {"dest": final_dst, "cat": info['cat']}
# Update DB # Update DB
cursor.execute("DELETE FROM staging_area WHERE original_path = ?", (file_path,)) cursor.execute("DELETE FROM staging_area WHERE original_path = ?", (file_path,))
cursor.execute("INSERT OR REPLACE INTO processed_log VALUES (?, ?, ?)", cursor.execute("INSERT OR REPLACE INTO processed_log VALUES (?, ?, ?)",
(file_path, info['cat'], operation)) (file_path, info['cat'], operation))
# --- CASE B: Cleanup --- # --- CASE B: Cleanup ---
elif cleanup_mode != "Keep": elif cleanup_mode != "Keep":
if cleanup_mode == "Move to Unused": if cleanup_mode == "Move to Unused":
unused_dir = os.path.join(os.path.dirname(file_path), "unused") unused_dir = os.path.join(os.path.dirname(file_path), "unused")
os.makedirs(unused_dir, exist_ok=True) os.makedirs(unused_dir, exist_ok=True)
dest_unused = os.path.join(unused_dir, os.path.basename(file_path)) dest_unused = os.path.join(unused_dir, os.path.basename(file_path))
shutil.move(file_path, dest_unused) shutil.move(file_path, dest_unused)
SorterEngine.fix_permissions(dest_unused) # Fix here too SorterEngine.fix_permissions(dest_unused) # Fix here too
elif cleanup_mode == "Delete": elif cleanup_mode == "Delete":
os.remove(file_path) os.remove(file_path)
conn.commit() conn.commit()
conn.close() conn.close()
return committed
@staticmethod @staticmethod
def rename_category(old_name, new_name): def rename_category(old_name, new_name):
@@ -708,4 +874,386 @@ class SorterEngine:
) )
result = {row[0]: {"cat": row[1], "index": row[2]} for row in cursor.fetchall()} result = {row[0]: {"cat": row[1], "index": row[2]} for row in cursor.fetchall()}
conn.close() conn.close()
return result return result
# --- 8. CAPTION SETTINGS & PROMPTS ---
@staticmethod
def get_caption_settings(profile):
"""Get caption API settings for a profile."""
conn = sqlite3.connect(SorterEngine.DB_PATH)
cursor = conn.cursor()
# Ensure table exists
cursor.execute('''CREATE TABLE IF NOT EXISTS caption_settings
(profile TEXT PRIMARY KEY,
api_endpoint TEXT DEFAULT 'http://localhost:8080/v1/chat/completions',
model_name TEXT DEFAULT 'local-model',
max_tokens INTEGER DEFAULT 300,
temperature REAL DEFAULT 0.7,
timeout_seconds INTEGER DEFAULT 60,
batch_size INTEGER DEFAULT 4)''')
cursor.execute("SELECT * FROM caption_settings WHERE profile = ?", (profile,))
row = cursor.fetchone()
conn.close()
if row:
return {
"profile": row[0],
"api_endpoint": row[1],
"model_name": row[2],
"max_tokens": row[3],
"temperature": row[4],
"timeout_seconds": row[5],
"batch_size": row[6]
}
else:
# Return defaults
return {
"profile": profile,
"api_endpoint": "http://localhost:8080/v1/chat/completions",
"model_name": "local-model",
"max_tokens": 300,
"temperature": 0.7,
"timeout_seconds": 60,
"batch_size": 4
}
@staticmethod
def save_caption_settings(profile, api_endpoint=None, model_name=None, max_tokens=None,
temperature=None, timeout_seconds=None, batch_size=None):
"""Save caption API settings for a profile."""
conn = sqlite3.connect(SorterEngine.DB_PATH)
cursor = conn.cursor()
# Ensure table exists
cursor.execute('''CREATE TABLE IF NOT EXISTS caption_settings
(profile TEXT PRIMARY KEY,
api_endpoint TEXT DEFAULT 'http://localhost:8080/v1/chat/completions',
model_name TEXT DEFAULT 'local-model',
max_tokens INTEGER DEFAULT 300,
temperature REAL DEFAULT 0.7,
timeout_seconds INTEGER DEFAULT 60,
batch_size INTEGER DEFAULT 4)''')
# Get existing values
cursor.execute("SELECT * FROM caption_settings WHERE profile = ?", (profile,))
row = cursor.fetchone()
if not row:
row = (profile, "http://localhost:8080/v1/chat/completions", "local-model", 300, 0.7, 60, 4)
new_values = (
profile,
api_endpoint if api_endpoint is not None else row[1],
model_name if model_name is not None else row[2],
max_tokens if max_tokens is not None else row[3],
temperature if temperature is not None else row[4],
timeout_seconds if timeout_seconds is not None else row[5],
batch_size if batch_size is not None else row[6]
)
cursor.execute("INSERT OR REPLACE INTO caption_settings VALUES (?, ?, ?, ?, ?, ?, ?)", new_values)
conn.commit()
conn.close()
@staticmethod
def get_category_prompt(profile, category):
"""Get prompt template for a category."""
conn = sqlite3.connect(SorterEngine.DB_PATH)
cursor = conn.cursor()
cursor.execute('''CREATE TABLE IF NOT EXISTS category_prompts
(profile TEXT, category TEXT, prompt_template TEXT,
PRIMARY KEY (profile, category))''')
cursor.execute(
"SELECT prompt_template FROM category_prompts WHERE profile = ? AND category = ?",
(profile, category)
)
row = cursor.fetchone()
conn.close()
if row and row[0]:
return row[0]
else:
# Default prompt
return "Describe this image in detail for training purposes. Include subjects, actions, setting, colors, and composition."
@staticmethod
def save_category_prompt(profile, category, prompt):
"""Save prompt template for a category."""
conn = sqlite3.connect(SorterEngine.DB_PATH)
cursor = conn.cursor()
cursor.execute('''CREATE TABLE IF NOT EXISTS category_prompts
(profile TEXT, category TEXT, prompt_template TEXT,
PRIMARY KEY (profile, category))''')
cursor.execute(
"INSERT OR REPLACE INTO category_prompts VALUES (?, ?, ?)",
(profile, category, prompt)
)
conn.commit()
conn.close()
@staticmethod
def get_all_category_prompts(profile):
"""Get all prompt templates for a profile."""
conn = sqlite3.connect(SorterEngine.DB_PATH)
cursor = conn.cursor()
cursor.execute('''CREATE TABLE IF NOT EXISTS category_prompts
(profile TEXT, category TEXT, prompt_template TEXT,
PRIMARY KEY (profile, category))''')
cursor.execute(
"SELECT category, prompt_template FROM category_prompts WHERE profile = ?",
(profile,)
)
result = {row[0]: row[1] for row in cursor.fetchall()}
conn.close()
return result
# --- 9. CAPTION STORAGE ---
@staticmethod
def save_caption(image_path, caption, model):
"""Save a generated caption to the database."""
conn = sqlite3.connect(SorterEngine.DB_PATH)
cursor = conn.cursor()
cursor.execute('''CREATE TABLE IF NOT EXISTS image_captions
(image_path TEXT PRIMARY KEY, caption TEXT, model TEXT,
generated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
cursor.execute(
"INSERT OR REPLACE INTO image_captions VALUES (?, ?, ?, ?)",
(image_path, caption, model, datetime.now().isoformat())
)
conn.commit()
conn.close()
@staticmethod
def get_caption(image_path):
"""Get caption for an image."""
conn = sqlite3.connect(SorterEngine.DB_PATH)
cursor = conn.cursor()
cursor.execute('''CREATE TABLE IF NOT EXISTS image_captions
(image_path TEXT PRIMARY KEY, caption TEXT, model TEXT,
generated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
cursor.execute(
"SELECT caption, model, generated_at FROM image_captions WHERE image_path = ?",
(image_path,)
)
row = cursor.fetchone()
conn.close()
if row:
return {"caption": row[0], "model": row[1], "generated_at": row[2]}
return None
@staticmethod
def get_captions_batch(image_paths):
"""Get captions for multiple images."""
if not image_paths:
return {}
conn = sqlite3.connect(SorterEngine.DB_PATH)
cursor = conn.cursor()
cursor.execute('''CREATE TABLE IF NOT EXISTS image_captions
(image_path TEXT PRIMARY KEY, caption TEXT, model TEXT,
generated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
placeholders = ','.join('?' * len(image_paths))
cursor.execute(
f"SELECT image_path, caption, model, generated_at FROM image_captions WHERE image_path IN ({placeholders})",
image_paths
)
result = {row[0]: {"caption": row[1], "model": row[2], "generated_at": row[3]} for row in cursor.fetchall()}
conn.close()
return result
@staticmethod
def delete_caption(image_path):
"""Delete caption for an image."""
conn = sqlite3.connect(SorterEngine.DB_PATH)
cursor = conn.cursor()
cursor.execute("DELETE FROM image_captions WHERE image_path = ?", (image_path,))
conn.commit()
conn.close()
@staticmethod
def get_all_caption_paths():
"""Get set of all image paths that have captions."""
conn = sqlite3.connect(SorterEngine.DB_PATH)
cursor = conn.cursor()
cursor.execute('''CREATE TABLE IF NOT EXISTS image_captions
(image_path TEXT PRIMARY KEY, caption TEXT, model TEXT,
generated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)''')
cursor.execute("SELECT image_path FROM image_captions")
result = {row[0] for row in cursor.fetchall()}
conn.close()
return result
# --- 10. VLLM API CAPTIONING ---
@staticmethod
def caption_image_vllm(image_path, prompt, settings):
"""
Generate caption for an image using VLLM API.
Args:
image_path: Path to the image file
prompt: Text prompt for captioning
settings: Dict with api_endpoint, model_name, max_tokens, temperature, timeout_seconds
Returns:
Tuple of (caption_text, error_message). If successful, error is None.
"""
try:
# Read and encode image
with open(image_path, 'rb') as f:
img_bytes = f.read()
b64_image = base64.b64encode(img_bytes).decode('utf-8')
# Determine MIME type
ext = os.path.splitext(image_path)[1].lower()
mime_types = {
'.jpg': 'image/jpeg',
'.jpeg': 'image/jpeg',
'.png': 'image/png',
'.webp': 'image/webp',
'.bmp': 'image/bmp',
'.tiff': 'image/tiff'
}
mime_type = mime_types.get(ext, 'image/jpeg')
# Build request payload (OpenAI-compatible format)
payload = {
"model": settings.get('model_name', 'local-model'),
"messages": [{
"role": "user",
"content": [
{"type": "text", "text": prompt},
{"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{b64_image}"}}
]
}],
"max_tokens": settings.get('max_tokens', 300),
"temperature": settings.get('temperature', 0.7)
}
# Make API request
response = requests.post(
settings.get('api_endpoint', 'http://localhost:8080/v1/chat/completions'),
json=payload,
timeout=settings.get('timeout_seconds', 60)
)
response.raise_for_status()
result = response.json()
caption = result['choices'][0]['message']['content']
return caption.strip(), None
except requests.Timeout:
return None, f"API timeout after {settings.get('timeout_seconds', 60)}s"
except requests.RequestException as e:
return None, f"API error: {str(e)}"
except KeyError as e:
return None, f"Invalid API response: missing {str(e)}"
except Exception as e:
return None, f"Error: {str(e)}"
@staticmethod
def caption_batch_vllm(image_paths, get_prompt_fn, settings, progress_cb=None):
"""
Caption multiple images using VLLM API.
Args:
image_paths: List of (image_path, category) tuples
get_prompt_fn: Function(category) -> prompt string
settings: Caption settings dict
progress_cb: Optional callback(current, total, status_msg) for progress updates
Returns:
Dict with results: {"success": count, "failed": count, "captions": {path: caption}}
"""
results = {"success": 0, "failed": 0, "captions": {}, "errors": {}}
total = len(image_paths)
for i, (image_path, category) in enumerate(image_paths):
if progress_cb:
progress_cb(i, total, f"Captioning {os.path.basename(image_path)}...")
prompt = get_prompt_fn(category)
caption, error = SorterEngine.caption_image_vllm(image_path, prompt, settings)
if caption:
# Save to database
SorterEngine.save_caption(image_path, caption, settings.get('model_name', 'local-model'))
results["captions"][image_path] = caption
results["success"] += 1
else:
# Store error
error_caption = f"[ERROR] {error}"
SorterEngine.save_caption(image_path, error_caption, settings.get('model_name', 'local-model'))
results["errors"][image_path] = error
results["failed"] += 1
if progress_cb:
progress_cb(total, total, "Complete!")
return results
@staticmethod
def write_caption_sidecar(image_path, caption):
"""
Write caption to a .txt sidecar file next to the image.
Args:
image_path: Path to the image file
caption: Caption text to write
Returns:
Path to sidecar file, or None on error
"""
try:
# Create sidecar path (same name, .txt extension)
base_path = os.path.splitext(image_path)[0]
sidecar_path = f"{base_path}.txt"
with open(sidecar_path, 'w', encoding='utf-8') as f:
f.write(caption)
# Fix permissions
SorterEngine.fix_permissions(sidecar_path)
return sidecar_path
except Exception as e:
print(f"Warning: Could not write sidecar for {image_path}: {e}")
return None
@staticmethod
def read_caption_sidecar(image_path):
"""
Read caption from a .txt sidecar file if it exists.
Args:
image_path: Path to the image file
Returns:
Caption text or None if no sidecar exists
"""
try:
base_path = os.path.splitext(image_path)[0]
sidecar_path = f"{base_path}.txt"
if os.path.exists(sidecar_path):
with open(sidecar_path, 'r', encoding='utf-8') as f:
return f.read().strip()
except Exception:
pass
return None

File diff suppressed because it is too large Load Diff

View File

@@ -1,3 +1,4 @@
streamlit streamlit
Pillow Pillow
nicegui nicegui
requests

33
start.sh Normal file → Executable file
View File

@@ -1,18 +1,25 @@
#!/bin/bash #!/bin/bash
# 1. Navigate to app directory # NiceSorter - Start Script
cd /app # Runs the NiceGUI gallery interface on port 8080
# 2. Install dependencies (Including NiceGUI if missing) set -e
# This checks your requirements.txt AND ensures nicegui is present
pip install --no-cache-dir -r requirements.txt
# 3. Start NiceGUI in the Background (&) # Navigate to app directory if running in container
# This runs silently while the script continues if [ -d "/app" ]; then
echo "🚀 Starting NiceGUI on Port 8080..." cd /app
python3 gallery_app.py & fi
# 4. Start Streamlit in the Foreground # Install dependencies if requirements.txt exists
# This keeps the container running if [ -f "requirements.txt" ]; then
echo "🚀 Starting Streamlit on Port 8501..." echo "📦 Installing dependencies..."
streamlit run app.py --server.port=8501 --server.address=0.0.0.0 pip install --no-cache-dir -q -r requirements.txt
fi
# Initialize database
echo "🗄️ Initializing database..."
python3 -c "from engine import SorterEngine; SorterEngine.init_db()"
# Start NiceGUI
echo "🚀 Starting NiceSorter on http://0.0.0.0:8080"
exec python3 gallery_app.py