Create app.py

app.py

@@ -0,0 +1,398 @@
+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
+)