app_old.py

"""Photographer's Archive — Gradio app entry point."""

import json
import logging
import os
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):
    if not uploaded_files:
        yield "Please upload at least one image or a folder of images.", gr.update(), gr.update()
        return

    # Determine collection name
    final_collection = new_collection_name.strip() if is_new_collection else collection_name
    if not final_collection:
        final_collection = "General"

    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 "No supported images found in upload (jpg, jpeg, png, webp, tiff).", gr.update(), gr.update()
        return

    yield f"Staging images to collection '{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"[{processed}/{total}] ({pct}%) {msg}", gr.update(), gr.update()
            else:
                yield msg, gr.update(), gr.update()
        
        rows, _ = load_caption_browser()
        cols = gr.Dropdown(choices=get_all_collections(), value=final_collection)
        yield f"✅ Done. Ingested to '{final_collection}'. Store has {entry_count()} images.", rows, cols
    except ValueError as e:
        yield f"Error: {e}", gr.update(), gr.update()


def _parse_meta(raw: str) -> dict | None:
    """Try to parse raw caption as JSON, with comma-fix fallback."""
    import re
    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 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(s) in store."


def run_search(query: str, collection: str = "All"):
    if not query or not query.strip():
        return [], "Please enter a search query."
    if entry_count() == 0:
        return [], "No images indexed yet. Upload and ingest some photos first."
    
    col_filter = None if collection == "All" else collection
    results = search(query.strip(), collection=col_filter)
    
    if not results:
        return [], f"No matching photos found in {collection} (threshold: {MIN_RELEVANCE})."
    gallery_items = [
        (r["path"], f"Score: {r['score']:.2f} — {r['caption'][:120]}…")
        for r in results
    ]
    return gallery_items, f"{len(results)} result(s) found."


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 images in collection '{collection_name}'."
    
    gallery_items = [(path, os.path.basename(path)) for path in entries.keys()]
    return gallery_items, f"{len(entries)} image(s) in '{collection_name}'."


ABOUT_TEXT = """
## ShutterSearch

A local-first photo search tool powered by **MiniCPM-V-4.6** (≤7B VLM).

**How it works:**
1. Upload photos using the Ingest page — select an existing collection or create a new one.
   The VLM runs on Modal's GPU and generates rich captions. Captions are cached locally.
2. Use the Search page to find photos by content, mood, or composition using natural language.
3. Browse your organized archive in the Collections page.

Built for the Hugging Face [Build Small](https://huggingface.co/build-small) hackathon.
"""


# --- UI Layout ---
with gr.Blocks() as demo:
    # Navigation state
    current_page = gr.State("home")

    with gr.Sidebar():
        gr.Markdown("# 📷 ShutterSearch")
        gr.Markdown("Modern Photo Archive")
        nav_home = gr.Button("🏠 Home", variant="ghost", size="lg")
        nav_ingest = gr.Button("📥 Ingest", variant="ghost", size="lg")
        nav_search = gr.Button("🔍 Search", variant="ghost", size="lg")
        nav_collections = gr.Button("📁 Collections", variant="ghost", size="lg")
        nav_captions = gr.Button("📝 Captions", variant="ghost", size="lg")
        nav_about = gr.Button("ℹ️ About", variant="ghost", size="lg")
        
        gr.HTML("<hr>")
        stats = gr.Label(value=f"Total Photos: {entry_count()}", label="Archive Stats")

    # --- Pages ---
    with gr.Column() as home_page:
        with gr.Row():
            with gr.Column(scale=2):
                gr.Markdown("""
                # Welcome to ShutterSearch
                ### Your intelligent local-first photo archive.
                
                ShutterSearch uses state-of-the-art vision models to understand your photography. 
                Keep your photos organized in collections and find them instantly with natural language search.
                
                - **Semantic Search**: Find "sunset on a beach" even if you didn't tag it.
                - **Automatic Captioning**: Powered by MiniCPM-V-4.6.
                - **Privacy First**: Everything runs on your terms.
                """)
                start_btn = gr.Button("Start Ingesting", variant="primary", size="lg")
            with gr.Column(scale=1):
                # Placeholder for random image
                gr.Image("https://images.unsplash.com/photo-1542038784456-1ea8e935640e?q=80&w=1000&auto=format&fit=crop", 
                         label="Photography", show_label=False, interactive=False)

    with gr.Column(visible=False) as ingest_page:
        gr.Markdown("# 📥 Ingest Photos")
        with gr.Row():
            with gr.Column(scale=2):
                upload = gr.File(
                    label="Upload Images",
                    file_count="multiple",
                    file_types=[".jpg", ".jpeg", ".png", ".webp", ".tiff"],
                    height=300,
                )
            with gr.Column(scale=1):
                gr.Markdown("### Collection Settings")
                coll_dropdown = gr.Dropdown(
                    choices=get_all_collections(), 
                    label="Select Existing Collection",
                    value="General"
                )
                use_new_coll = gr.Checkbox(label="Create New Collection", value=False)
                new_coll_name = gr.Textbox(label="New Collection Name", visible=False)
                
                ingest_btn = gr.Button("🚀 Start Ingestion", variant="primary", size="lg")
        
        ingest_status = gr.Textbox(label="Status", interactive=False, lines=5)

    with gr.Column(visible=False) as search_page:
        gr.Markdown("# 🔍 AI Search")
        with gr.Row():
            search_query = gr.Textbox(
                placeholder="Describe what you're looking for... (e.g. 'moody forest portrait')",
                label="Search Query",
                scale=4
            )
            search_col_filter = gr.Dropdown(
                choices=["All"] + get_all_collections(),
                value="All",
                label="Filter by Collection",
                scale=1
            )
            search_btn = gr.Button("Search", variant="primary", scale=1)
        
        search_status_msg = gr.Markdown("Enter a query to start searching.")
        search_gallery = gr.Gallery(label="Search Results", columns=4, height="auto")

    with gr.Column(visible=False) as collections_page:
        gr.Markdown("# 📁 Collections")
        with gr.Row():
            view_coll_selector = gr.Dropdown(
                choices=get_all_collections(),
                value=get_all_collections()[0] if get_all_collections() else "General",
                label="Select Collection to Browse",
                scale=4
            )
            refresh_coll_btn = gr.Button("🔄 Refresh", scale=1)
        
        coll_status = gr.Markdown("Browse your organized photos.")
        coll_gallery = gr.Gallery(label="Collection Photos", columns=5, height="auto")

    with gr.Column(visible=False) as captions_page:
        gr.Markdown("# 📝 Caption Browser")
        with gr.Row():
            refresh_cap_btn = gr.Button("🔄 Refresh Data", variant="secondary")
            cap_status = gr.Textbox(interactive=False, show_label=False, scale=4)
            
        caption_table = gr.Dataframe(
            headers=["File", "Summary", "Attire", "Tags"],
            datatype=["str", "str", "str", "str"],
            wrap=True,
            interactive=False,
            column_widths=["15%", "45%", "20%", "20%"],
        )

    with gr.Column(visible=False) as about_page:
        gr.Markdown("# ℹ️ About ShutterSearch")
        gr.Markdown(ABOUT_TEXT)

    # --- Navigation Logic ---
    pages = [home_page, ingest_page, search_page, collections_page, captions_page, about_page]
    
    def switch_page(page_name):
        updates = [gr.update(visible=(name == page_name)) for name in ["home", "ingest", "search", "collections", "captions", "about"]]
        return updates

    nav_home.click(fn=lambda: switch_page("home"), outputs=pages)
    nav_ingest.click(fn=lambda: switch_page("ingest"), outputs=pages)
    nav_search.click(fn=lambda: switch_page("search") + [gr.update(choices=["All"] + get_all_collections())], outputs=pages + [search_col_filter])
    nav_collections.click(fn=lambda: switch_page("collections") + [gr.update(choices=get_all_collections())], outputs=pages + [view_coll_selector])
    nav_captions.click(fn=lambda: switch_page("captions"), outputs=pages)
    nav_about.click(fn=lambda: switch_page("about"), outputs=pages)
    
    start_btn.click(fn=lambda: switch_page("ingest"), outputs=pages)

    # --- Page Interactions ---
    use_new_coll.change(fn=lambda x: gr.update(visible=x), inputs=use_new_coll, outputs=new_coll_name)
    
    ingest_btn.click(
        fn=run_ingest, 
        inputs=[upload, coll_dropdown, use_new_coll, new_coll_name], 
        outputs=[ingest_status, caption_table, coll_dropdown]
    ).then(
        fn=lambda: f"Total Photos: {entry_count()}", outputs=stats
    )

    search_btn.click(fn=run_search, inputs=[search_query, search_col_filter], outputs=[search_gallery, search_status_msg])
    search_query.submit(fn=run_search, inputs=[search_query, search_col_filter], outputs=[search_gallery, search_status_msg])

    view_coll_selector.change(fn=load_collections_view, inputs=view_coll_selector, outputs=[coll_gallery, coll_status])
    refresh_coll_btn.click(fn=load_collections_view, inputs=view_coll_selector, outputs=[coll_gallery, coll_status])

    refresh_cap_btn.click(fn=load_caption_browser, outputs=[caption_table, cap_status])
    
    demo.load(fn=load_caption_browser, outputs=[caption_table, cap_status])
    demo.load(fn=lambda: load_collections_view(get_all_collections()[0] if get_all_collections() else "General"), outputs=[coll_gallery, coll_status])

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Soft(primary_hue="green", secondary_hue="emerald"))