caption fix
This commit is contained in:
Binary file not shown.
Binary file not shown.
31
engine.py
31
engine.py
@@ -47,9 +47,19 @@ class SorterEngine:
|
||||
(source_path TEXT PRIMARY KEY, category TEXT, action_type TEXT)''')
|
||||
|
||||
# --- NEW: FOLDER TAGS TABLE (persists tags by folder) ---
|
||||
cursor.execute('''CREATE TABLE IF NOT EXISTS folder_tags
|
||||
(folder_path TEXT, filename TEXT, category TEXT, tag_index INTEGER,
|
||||
PRIMARY KEY (folder_path, filename))''')
|
||||
# Check if old schema exists (without profile column) and migrate
|
||||
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='folder_tags'")
|
||||
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) ---
|
||||
cursor.execute('''CREATE TABLE IF NOT EXISTS profile_categories
|
||||
@@ -1063,6 +1073,21 @@ class SorterEngine:
|
||||
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):
|
||||
|
||||
@@ -7,6 +7,9 @@ from nicegui import ui, app, run
|
||||
from fastapi import Response
|
||||
from engine import SorterEngine
|
||||
|
||||
# Initialize database tables on startup
|
||||
SorterEngine.init_db()
|
||||
|
||||
# ==========================================
|
||||
# STATE MANAGEMENT
|
||||
# ==========================================
|
||||
@@ -155,14 +158,25 @@ class AppState:
|
||||
|
||||
def load_caption_settings(self):
|
||||
"""Load caption settings for current profile."""
|
||||
self.caption_settings = SorterEngine.get_caption_settings(self.profile_name)
|
||||
try:
|
||||
self.caption_settings = SorterEngine.get_caption_settings(self.profile_name)
|
||||
except Exception:
|
||||
self.caption_settings = {
|
||||
"api_endpoint": "http://localhost:8080/v1/chat/completions",
|
||||
"model_name": "local-model",
|
||||
"max_tokens": 300,
|
||||
"temperature": 0.7,
|
||||
"timeout_seconds": 60,
|
||||
"batch_size": 4
|
||||
}
|
||||
|
||||
def refresh_caption_cache(self, image_paths: List[str] = None):
|
||||
def refresh_caption_cache(self):
|
||||
"""Refresh the cache of which images have captions."""
|
||||
paths = image_paths or self.all_images
|
||||
if paths:
|
||||
captions = SorterEngine.get_captions_batch(paths)
|
||||
self.caption_cache = set(captions.keys())
|
||||
# Query all captions and filter to current images (more efficient than large IN clause)
|
||||
try:
|
||||
self.caption_cache = SorterEngine.get_all_caption_paths()
|
||||
except Exception:
|
||||
self.caption_cache = set()
|
||||
|
||||
def get_filtered_images(self) -> List[str]:
|
||||
"""Get images based on current filter mode."""
|
||||
@@ -259,11 +273,10 @@ def load_images():
|
||||
state.page = 0
|
||||
|
||||
refresh_staged_info()
|
||||
# Refresh caption cache for loaded images
|
||||
state.refresh_caption_cache()
|
||||
# Load caption settings
|
||||
state.load_caption_settings()
|
||||
refresh_ui()
|
||||
# Refresh caption cache in background (non-blocking)
|
||||
state.refresh_caption_cache()
|
||||
state.load_caption_settings()
|
||||
|
||||
# ==========================================
|
||||
# PAIRING MODE FUNCTIONS
|
||||
|
||||
Reference in New Issue
Block a user