ShutterSearch / logic.py
SwikarG's picture
Merge github code to hugging face (#1)
5f39285
Raw
History Blame Contribute Delete
14.8 kB
# import json
# import logging
# import os
# import re
# import shutil
# import tempfile
# import gradio as gr
# from caption_store import all_entries, entry_count, get_all_collections, get_entries_by_collection
# from ingest import ingest_folder
# from search import MIN_RELEVANCE, search
# logging.basicConfig(level=logging.INFO)
# IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".tiff"}
# STAGING_DIR = os.path.join(tempfile.gettempdir(), "photographers_archive_uploads")
# os.makedirs(STAGING_DIR, exist_ok=True)
# def run_ingest(uploaded_files, collection_name, is_new_collection, new_collection_name):
# """
# Clears the staging directory, moves uploaded images, validates the destination
# collection, and runs the vision-model ingestion process.
# """
# if not uploaded_files:
# yield "⚠️ Please upload at least one image to begin.", gr.update(), gr.update()
# return
# # 1. Determine and validate collection selection context
# if is_new_collection:
# final_collection = new_collection_name.strip()
# if not final_collection:
# yield "⚠️ Ingestion halted: Please specify a valid name for the new collection.", gr.update(), gr.update()
# return
# else:
# final_collection = collection_name
# if not final_collection:
# final_collection = "General"
# # 2. Housekeep staging directory (clear residuals from prior sessions)
# try:
# for filename in os.listdir(STAGING_DIR):
# file_path = os.path.join(STAGING_DIR, filename)
# if os.path.isfile(file_path) or os.path.islink(file_path):
# os.unlink(file_path)
# elif os.path.isdir(file_path):
# shutil.rmtree(file_path)
# except Exception as e:
# logging.warning(f"Could not clear staging directory fully: {e}")
# # 3. Stage files
# staged = []
# for file_path in uploaded_files:
# ext = os.path.splitext(file_path)[-1].lower()
# if ext in IMAGE_EXTENSIONS:
# dest = os.path.join(STAGING_DIR, os.path.basename(file_path))
# shutil.copy2(file_path, dest)
# staged.append(dest)
# if not staged:
# yield "⚠️ Format error: No supported images (jpg, jpeg, png, webp, tiff) found.", gr.update(), gr.update()
# return
# yield f"Staging completed. Initializing pipeline for: '{final_collection}'...", gr.update(), gr.update()
# # 4. Ingest and track progress
# try:
# for processed, total, msg in ingest_folder(STAGING_DIR, collection=final_collection):
# if total > 0:
# pct = int(processed / total * 100)
# yield f">> Progress: [{processed}/{total}] ({pct}%) - {msg}", gr.update(), gr.update()
# else:
# yield f">> {msg}", gr.update(), gr.update()
# # 5. Fetch updated data and regenerate dropdown choices safely
# rows, _ = load_caption_browser()
# choices = get_all_collections()
# if not choices:
# choices = ["General"]
# dropdown_val = final_collection if final_collection in choices else choices[0]
# cols = gr.Dropdown(choices=choices, value=dropdown_val)
# yield f"System Log: Done. Ingested to '{final_collection}'. Store has {entry_count()} images.", rows, cols
# except ValueError as e:
# yield f"Pipeline Failure: {e}", gr.update(), gr.update()
# def _parse_meta(raw: str) -> dict | None:
# """Attempts to parse raw metadata caption as JSON with syntax fallback support."""
# try:
# return json.loads(raw)
# except (json.JSONDecodeError, TypeError):
# pass
# try:
# # Fallback fix for missing trailing commas in lists
# fixed = re.sub(r'"\s*\n(\s*")', r'",\n\1', raw)
# return json.loads(fixed)
# except (json.JSONDecodeError, TypeError):
# return None
# def load_caption_browser():
# """Loads all caption entries structured for tabular visualization."""
# entries = all_entries()
# if not entries:
# return [], "No captions index exists yet."
# rows = []
# for path, data in entries.items():
# raw = data["caption"]
# meta = _parse_meta(raw)
# if meta:
# summary = meta.get("summary") or "—"
# subj = meta.get("subjects", {})
# attire = ", ".join(subj.get("attire", [])) or "—"
# tags = ", ".join(meta.get("search_tags", [])) or "—"
# else:
# summary = raw[:300] if raw else "—"
# attire = "—"
# tags = "—"
# rows.append([os.path.basename(path), summary, attire, tags])
# return rows, f"{len(rows)} caption record(s) loaded."
# def run_search(query: str, collection: str = "All"):
# """Searches indexed captions using natural language match thresholding."""
# if not query or not query.strip():
# return [], "Please enter a valid search parameter."
# if entry_count() == 0:
# return [], "No images indexed yet. Please ingest photos first."
# col_filter = None if collection == "All" else collection
# results = search(query.strip(), collection=col_filter)
# if not results:
# return [], f"Zero matches in database for target {collection} (Threshold constraint: {MIN_RELEVANCE})."
# gallery_items = [
# (r["path"], f"Confidence: {r['score']:.2f} | {r['caption'][:120]}…")
# for r in results
# ]
# return gallery_items, f"{len(results)} query match(es) located."
# def load_collections_view(collection_name):
# """Fetches list of path references sorted within specific collection targets."""
# if not collection_name or collection_name == "All":
# entries = all_entries()
# else:
# entries = get_entries_by_collection(collection_name)
# if not entries:
# return [], f"No stored assets in target collection: '{collection_name}'."
# gallery_items = [(path, os.path.basename(path)) for path in entries.keys()]
# return gallery_items, f"Found {len(entries)} file reference(s) within '{collection_name}'."
# def update_collections_dropdown():
# """Returns a safe state package configuration payload for collection selectors."""
# choices = get_all_collections()
# if not choices:
# choices = ["General"]
# val = "General" if "General" in choices else choices[0]
# return gr.update(choices=choices, value=val)
# def update_search_dropdown():
# """Returns a safe state package configuration payload for search filter dropdowns."""
# choices = ["All"] + get_all_collections()
# return gr.update(choices=choices, value="All")
import hashlib
import json
import logging
import os
import re
import shutil
import tempfile
import zipfile
import gradio as gr
from PIL import Image
from caption_store import all_entries, entry_count, get_all_collections, get_entries_by_collection
from ingest import ingest_folder
from search import MIN_RELEVANCE, search
logging.basicConfig(level=logging.INFO)
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".tiff"}
# STAGING_DIR = os.path.join(tempfile.gettempdir(), "photographers_archive_uploads")
STAGING_DIR = "./photographers_archive_uploads"
os.makedirs(STAGING_DIR, exist_ok=True)
THUMBNAIL_DIR = os.path.join(tempfile.gettempdir(), "photographers_archive_thumbnails")
os.makedirs(STAGING_DIR, exist_ok=True)
os.makedirs(THUMBNAIL_DIR, exist_ok=True)
def get_thumbnail_path(original_path):
"""Generates and returns a cached lightweight WebP thumbnail path."""
path_hash = hashlib.md5(original_path.encode('utf-8')).hexdigest()
thumb_path = os.path.join(THUMBNAIL_DIR, f"{path_hash}.webp")
if os.path.exists(thumb_path):
return thumb_path
try:
with Image.open(original_path) as img:
img.thumbnail((300, 300))
img.save(thumb_path, "WEBP", quality=70)
return thumb_path
except Exception as e:
logging.warning(f"Could not render thumbnail for {original_path}: {e}")
return original_path
def run_ingest(uploaded_files, collection_name, is_new_collection, new_collection_name):
if not uploaded_files:
yield "⚠️ Please upload at least one image to begin.", gr.update(), gr.update()
return
if is_new_collection:
final_collection = new_collection_name.strip()
if not final_collection:
yield "⚠️ Ingestion halted: Please specify a valid name for the new collection.", gr.update(), gr.update()
return
else:
final_collection = collection_name
if not final_collection:
final_collection = "General"
try:
for filename in os.listdir(STAGING_DIR):
file_path = os.path.join(STAGING_DIR, filename)
if os.path.isfile(file_path) or os.path.islink(file_path):
os.unlink(file_path)
except Exception as e:
logging.warning(f"Could not clear staging directory fully: {e}")
staged = []
for file_path in uploaded_files:
ext = os.path.splitext(file_path)[-1].lower()
if ext in IMAGE_EXTENSIONS:
dest = os.path.join(STAGING_DIR, os.path.basename(file_path))
shutil.copy2(file_path, dest)
staged.append(dest)
if not staged:
yield "⚠️ Format error: No supported images (jpg, jpeg, png, webp, tiff) found.", gr.update(), gr.update()
return
yield f"Staging completed. Initializing pipeline for: '{final_collection}'...", gr.update(), gr.update()
try:
for processed, total, msg in ingest_folder(STAGING_DIR, collection=final_collection):
if total > 0:
pct = int(processed / total * 100)
yield f">> Progress: [{processed}/{total}] ({pct}%) - {msg}", gr.update(), gr.update()
else:
yield f">> {msg}", gr.update(), gr.update()
rows, _ = load_caption_browser()
choices = get_all_collections()
if not choices:
choices = ["General"]
dropdown_val = final_collection if final_collection in choices else choices[0]
cols = gr.Dropdown(choices=choices, value=dropdown_val)
yield f"System Log: Done. Ingested to '{final_collection}'. Store has {entry_count()} images.", rows, cols
except ValueError as e:
yield f"Pipeline Failure: {e}", gr.update(), gr.update()
def _parse_meta(raw: str) -> dict | None:
try:
return json.loads(raw)
except (json.JSONDecodeError, TypeError):
pass
try:
fixed = re.sub(r'"\s*\n(\s*")', r'",\n\1', raw)
return json.loads(fixed)
except (json.JSONDecodeError, TypeError):
return None
def load_caption_browser():
entries = all_entries()
if not entries:
return [], "No captions index exists yet."
rows = []
for path, data in entries.items():
raw = data["caption"]
meta = _parse_meta(raw)
if meta:
summary = meta.get("summary") or "—"
subj = meta.get("subjects", {})
attire = ", ".join(subj.get("attire", [])) or "—"
tags = ", ".join(meta.get("search_tags", [])) or "—"
else:
summary = raw[:300] if raw else "—"
attire = "—"
tags = "—"
rows.append([os.path.basename(path), summary, attire, tags])
return rows, f"{len(rows)} caption record(s) loaded."
def run_search(query: str, collection: str = "All"):
"""Performs search and outputs thumbnail images, absolute original files, and logs."""
if not query or not query.strip():
return [], [], "Please enter a valid search parameter."
if entry_count() == 0:
return [], [], "No images indexed yet. Please ingest photos first."
col_filter = None if collection == "All" else collection
results = search(query.strip(), collection=col_filter)
CUSTOM_THRESHOLD = 0.60
filtered_results = [r for r in results if r.get("score", 0) >= CUSTOM_THRESHOLD]
if not filtered_results:
return [], [], f"Zero matches in database for target {collection} (Threshold constraint: {MIN_RELEVANCE})."
original_paths = [r["path"] for r in filtered_results]
gallery_items = []
for r in filtered_results:
thumb = get_thumbnail_path(r["path"])
gallery_items.append((thumb, os.path.basename(r["path"])))
return gallery_items, original_paths, f"Found {len(filtered_results)} search matches."
def load_collections_view(collection_name):
if not collection_name or collection_name == "All":
entries = all_entries()
else:
entries = get_entries_by_collection(collection_name)
if not entries:
return [], [], f"No stored assets in target collection: '{collection_name}'."
original_paths = list(entries.keys())
gallery_items = []
for path in original_paths:
thumb = get_thumbnail_path(path)
gallery_items.append((thumb, os.path.basename(path)))
return gallery_items, original_paths, f"Found {len(entries)} image(s) within '{collection_name}'."
def zip_selected_files(selected_list):
if not selected_list:
return None, "⚠️ Downloader: Zero images selected."
try:
temp_zip = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
with zipfile.ZipFile(temp_zip.name, 'w', zipfile.ZIP_DEFLATED) as zipf:
for file_path in selected_list:
if os.path.exists(file_path):
zipf.write(file_path, os.path.basename(file_path))
return temp_zip.name, f"✅ Zip file ready with {len(selected_list)} source file(s)."
except Exception as e:
return None, f"⚠️ Compression failure: {e}"
def update_collections_dropdown():
choices = get_all_collections()
if not choices:
choices = ["General"]
val = "General" if "General" in choices else choices[0]
return gr.update(choices=choices, value=val)
def update_search_dropdown():
choices = ["All"] + get_all_collections()
return gr.update(choices=choices, value="All")