Delete src

src/streamlit_app.py

@@ -1,398 +0,0 @@
-import streamlit as st
-import io
-import base64
-import pandas as pd
-from PIL import Image
-from datetime import datetime
-import csv
-import json
-import os
-import requests
-
-# Optional PDF support via PyMuPDF
-try:
-    import fitz  # PyMuPDF
-    PDF_SUPPORT = True
-except ImportError:
-    PDF_SUPPORT = False
-
-# ---------------------------
-# Page config
-# ---------------------------
-st.set_page_config(
-    page_title="Curiosity AI Scans",
-    page_icon="🔍",
-    layout="wide",
-    initial_sidebar_state="expanded"
-)
-
-# ---------------------------
-# Helpers
-# ---------------------------
-def resize_image(image, max_size=1920):
-    w, h = image.size
-    if w > max_size or h > max_size:
-        if w > h:
-            nw = max_size
-            nh = int(h * (max_size / w))
-        else:
-            nh = max_size
-            nw = int(w * (max_size / h))
-        return image.resize((nw, nh), Image.LANCZOS)
-    return image
-
-def image_to_base64(image):
-    buf = io.BytesIO()
-    image.save(buf, format='JPEG')
-    return base64.b64encode(buf.getvalue()).decode('utf-8')
-
-def extract_structured_data(content, fields):
-    """Try to pull a JSON object for the requested fields out of model text."""
-    structured_data = {}
-    try:
-        # Fenced JSON
-        if "```json" in content and "```" in content.split("```json")[1]:
-            json_str = content.split("```json")[1].split("```")[0].strip()
-            structured_data.update(json.loads(json_str))
-        else:
-            # As a fallback, attempt to parse whole content if it looks like JSON
-            try:
-                maybe = json.loads(content)
-                if isinstance(maybe, dict):
-                    structured_data.update(maybe)
-            except Exception:
-                pass
-    except Exception:
-        pass
-    return structured_data
-
-# ---------------------------
-# OpenRouter client
-# ---------------------------
-OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")  # set this in Space Secrets
-
-def query_openrouter(prompt: str, image_base64: str, model_id: str) -> str:
-    if not OPENROUTER_API_KEY:
-        raise RuntimeError("Missing OPENROUTER_API_KEY. Add it in your Space → Settings → Variables & secrets.")
-
-    data_url = f"data:image/jpeg;base64,{image_base64}"
-
-    payload = {
-        "model": model_id,  # e.g., "google/gemma-3-4b-it"
-        "messages": [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": prompt},
-                    {"type": "image_url", "image_url": {"url": data_url}}
-                ]
-            }
-        ],
-        "max_tokens": 800
-    }
-
-    headers = {
-        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
-        "Content-Type": "application/json",
-        # Optional but recommended for attribution
-        "HTTP-Referer": st.secrets.get("SPACE_URL", "https://hf.space"),
-        "X-Title": "Curiosity AI Scans"
-    }
-
-    r = requests.post(
-        "https://openrouter.ai/api/v1/chat/completions",
-        headers=headers,
-        json=payload,
-        timeout=120
-    )
-    r.raise_for_status()
-    data = r.json()
-    return data["choices"][0]["message"]["content"]
-
-# ---------------------------
-# Core processing
-# ---------------------------
-def process_image(image, filename, fields=None, model=None):
-    img_base64 = image_to_base64(resize_image(image))
-
-    if fields is None:
-        prompt = "Describe this image in detail."
-        content = query_openrouter(prompt, img_base64, model)
-        return {'filename': filename, 'description': content}, content, None
-    else:
-        fields_str = ", ".join(fields)
-        prompt = (
-            "Extract the following fields from this image and return JSON only "
-            f"with these exact keys: {fields_str}. If a field is missing, use an empty string."
-        )
-        content = query_openrouter(prompt, img_base64, model)
-        structured_data = {'filename': filename}
-        parsed = extract_structured_data(content, fields)
-        if parsed:
-            structured_data.update(parsed)
-        return {'filename': filename, 'extraction': content}, content, structured_data
-
-def process_pdf(file_bytes, filename, fields=None, process_pages_separately=True, model=None):
-    """Rasterize PDF pages and run them through the same image path."""
-    if not PDF_SUPPORT:
-        yield None, None, None, filename, "PDF support requires PyMuPDF. Install pymupdf.", None
-        return
-
-    try:
-        pdf_document = fitz.open(stream=file_bytes, filetype="pdf")
-        page_count = len(pdf_document)
-
-        if process_pages_separately:
-            for page_num in range(page_count):
-                page = pdf_document[page_num]
-                pix = page.get_pixmap(matrix=fitz.Matrix(1.5, 1.5))
-                img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-                page_filename = f"{filename} (Page {page_num+1})"
-                result, content, structured_data = process_image(img, page_filename, fields, model)
-                yield page_num, page_count, img, page_filename, content, structured_data
-        else:
-            page = pdf_document[0]
-            pix = page.get_pixmap(matrix=fitz.Matrix(1.5, 1.5))
-            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-            result, content, structured_data = process_image(img, filename, fields, model)
-            yield 0, page_count, img, filename, content, structured_data
-
-    except Exception as e:
-        yield None, None, None, filename, f"Error processing PDF: {str(e)}", None
-
-def create_download_buttons(results, structured_results, extraction_mode):
-    st.header("Download Results")
-
-    # Simple CSV of descriptions or raw extraction
-    base_csv = io.StringIO()
-    base_writer = csv.writer(base_csv)
-    base_writer.writerow(['Filename', 'Description/Extraction'])
-    for r in results:
-        base_writer.writerow([r['filename'], r.get('description', r.get('extraction', ''))])
-
-    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
-    base_name = f"image_analysis_{ts}.csv"
-
-    st.success("All files processed.")
-    st.download_button(
-        label="Download Results (CSV)",
-        data=base_csv.getvalue(),
-        file_name=base_name,
-        mime="text/csv",
-        use_container_width=True
-    )
-
-    # Structured CSV if available
-    if extraction_mode == "Custom field extraction" and structured_results:
-        all_fields = set(['filename'])
-        for row in structured_results:
-            all_fields.update(row.keys())
-        headers = sorted(list(all_fields))
-        buff = io.StringIO()
-        w = csv.writer(buff)
-        w.writerow(headers)
-        for row in structured_results:
-            w.writerow([row.get(h, '') for h in headers])
-        st.download_button(
-            label="Download Structured Data (CSV)",
-            data=buff.getvalue(),
-            file_name=f"structured_data_{ts}.csv",
-            mime="text/csv",
-            use_container_width=True
-        )
-
-# ---------------------------
-# UI
-# ---------------------------
-st.title("Curiosity AI Scans")
-
-# Session state
-if 'results' not in st.session_state:
-    st.session_state.results = []
-if 'structured_results' not in st.session_state:
-    st.session_state.structured_results = []
-
-# Sidebar
-with st.sidebar:
-    st.header("Upload Files")
-    uploaded_files = st.file_uploader(
-        "Choose images or PDFs",
-        accept_multiple_files=True,
-        type=['png', 'jpg', 'jpeg', 'pdf']
-    )
-
-    st.header("Model Settings")
-    # OpenRouter model id for Gemma 3 4B Instruct (vision)
-    selected_model = st.selectbox(
-        "Choose vision model:",
-        ["google/gemma-3-4b-it"],
-        help="OpenRouter model id"
-    )
-
-    extraction_mode = "General description"
-    pdf_process_mode = "Process each page separately"
-    fields = None
-
-    if uploaded_files:
-        st.write(f"Uploaded {len(uploaded_files)} file(s)")
-
-        st.header("Data Extraction Options")
-        extraction_mode = st.radio(
-            "Choose extraction mode:",
-            ["General description", "Custom field extraction"]
-        )
-
-        if extraction_mode == "Custom field extraction":
-            custom_fields = st.text_area(
-                "Enter fields to extract (comma separated):",
-                value="Invoice number, Date, Company name, Total amount"
-            )
-            fields = [f.strip() for f in custom_fields.split(",") if f.strip()]
-
-            if any(file.name.lower().endswith('.pdf') for file in uploaded_files):
-                pdf_process_mode = st.radio(
-                    "How to process PDF files:",
-                    ["Process each page separately", "Process entire PDF as one document"]
-                )
-
-        process_button = st.button("Process Files", use_container_width=True)
-    else:
-        process_button = False
-        st.info("Upload images or PDFs to begin.")
-
-# Main processing
-if uploaded_files and process_button:
-    if not OPENROUTER_API_KEY:
-        st.error("OPENROUTER_API_KEY is not set. Add it in your Space → Settings → Variables & secrets.")
-    else:
-        st.header("Processing Results")
-        progress_bar = st.progress(0)
-        status_text = st.empty()
-
-        st.session_state.results = []
-        st.session_state.structured_results = []
-
-        # Count items to process
-        total_items = 0
-        for f in uploaded_files:
-            file_bytes = f.read()
-            f.seek(0)
-            if f.name.lower().endswith('.pdf') and PDF_SUPPORT:
-                if pdf_process_mode == "Process each page separately":
-                    try:
-                        pdf_document = fitz.open(stream=file_bytes, filetype="pdf")
-                        total_items += len(pdf_document)
-                    except Exception:
-                        total_items += 1
-                else:
-                    total_items += 1
-            else:
-                total_items += 1
-
-        processed_count = 0
-
-        # Process files
-        for f in uploaded_files:
-            file_bytes = f.read()
-            f.seek(0)
-
-            if f.name.lower().endswith('.pdf'):
-                if not PDF_SUPPORT:
-                    st.error("PDF support requires PyMuPDF. Add 'pymupdf' to requirements.txt.")
-                    processed_count += 1
-                    progress_bar.progress(processed_count / max(total_items, 1))
-                    continue
-
-                try:
-                    process_separately = pdf_process_mode == "Process each page separately"
-                    for page_info in process_pdf(file_bytes, f.name, fields, process_separately, selected_model):
-                        page_num, page_count, image, page_filename, content, structured_data = page_info
-                        if page_num is None:
-                            st.error(content)
-                            continue
-
-                        status_text.text(f"Processing {page_filename} ({page_num+1}/{page_count})")
-                        result = {'filename': page_filename, 'description': content}
-                        st.session_state.results.append(result)
-                        if structured_data and len(structured_data) > 1:
-                            st.session_state.structured_results.append(structured_data)
-
-                        st.subheader(page_filename)
-                        c1, c2 = st.columns([1, 2])
-                        with c1:
-                            st.image(image, width=250)
-                            if page_count > 1 and not process_separately:
-                                st.info(f"PDF has {page_count} pages. Showing first page only.")
-                        with c2:
-                            st.write(content)
-                            if structured_data and len(structured_data) > 1:
-                                st.success("Extracted structured data")
-                                st.json(structured_data)
-
-                        st.divider()
-                        processed_count += 1
-                        progress_bar.progress(min(processed_count / max(total_items, 1), 1.0))
-
-                except Exception as e:
-                    st.error(f"Error processing PDF {f.name}: {e}")
-                    processed_count += 1
-                    progress_bar.progress(min(processed_count / max(total_items, 1), 1.0))
-
-            else:
-                try:
-                    status_text.text(f"Processing image {f.name}")
-                    image = Image.open(f).convert("RGB")
-                    result, content, structured_data = process_image(image, f.name, fields, selected_model)
-                    st.session_state.results.append(result)
-                    if structured_data and len(structured_data) > 1:
-                        st.session_state.structured_results.append(structured_data)
-
-                    st.subheader(f"Image: {f.name}")
-                    c1, c2 = st.columns([1, 2])
-                    with c1:
-                        st.image(image, width=250)
-                    with c2:
-                        st.write(content)
-                        if structured_data and len(structured_data) > 1:
-                            st.success("Extracted structured data")
-                            st.json(structured_data)
-
-                    st.divider()
-
-                except Exception as e:
-                    st.error(f"Error processing image {f.name}: {e}")
-
-                processed_count += 1
-                progress_bar.progress(min(processed_count / max(total_items, 1), 1.0))
-
-        status_text.text("Processing complete.")
-
-        if st.session_state.results:
-            create_download_buttons(
-                st.session_state.results,
-                st.session_state.structured_results,
-                extraction_mode
-            )
-
-# Empty state
-if not uploaded_files:
-    st.info("Upload files using the sidebar to get started.")
-    st.write("""
-    How to use:
-    1) Upload one or more images or PDFs
-    2) Choose the OpenRouter vision model (Gemma 3 4B IT)
-    3) Pick description or custom field extraction
-    4) For PDFs, choose page-by-page or first page
-    5) Click Process Files
-    6) Review outputs and download CSVs
-    """)
-
-st.markdown("---")
-st.markdown(
-    """
-    <div style="text-align: center; margin-top: 12px; opacity: 0.7;">
-        Built for Hugging Face Spaces + OpenRouter
-    </div>
-    """,
-    unsafe_allow_html=True
-)