Spaces:

build-small-hackathon
/

ShutterSearch

Runtime error

File size: 14,791 Bytes

# 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")