Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,358 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+import os
+import re
+import io
+import json
 import streamlit as st
+from PIL import Image, ImageDraw, ImageFont
+from google import generativeai as genai # Renamed for clarity from 'google.genai'
+from google.generativeai.types import GenerationConfig, HarmCategory, HarmBlockThreshold
+from pdf2image import convert_from_bytes
+from typing import List, Dict, Any, Tuple, Optional
 
-"""
-# Welcome to Streamlit!
+# --- Constants ---
+MODEL_NAME = "gemini-2.5-pro-exp-03-25" # Using Gemini 1.5 Flash
+
+# Unified prompt to get bounding boxes and text content in one go
+UNIFIED_DETECTION_EXTRACTION_PROMPT = """\
+Analyze this document image. Identify text regions and extract the text from each region.
+
+Follow these rules for text regions:
+1. GROUP RELATED CONTENT:
+    - Full tables as SINGLE regions (including headers and all rows).
+    - Paragraphs as SINGLE rectangular blocks (multiple lines as one box).
+    - Keep text columns intact.
+    - Treat list items as single region if visually grouped.
+
+2. TEXT REGION REQUIREMENTS:
+    - Boundaries must tightly wrap text content.
+    - Include approximately 2% padding around text clusters, but ensure the box stays within image bounds.
+    - Exclude isolated decorative elements unless they contain text.
+    - Merge adjacent text fragments with ≤1% spacing into a single region.
+
+3. COORDINATE FORMAT:
+    - Normalized coordinates (0.0 to 1.0) with 3 decimal places.
+    - Format: [xmin, ymin, xmax, ymax].
+    - Order: Top-to-bottom, then left-to-right.
+
+4. SPECIAL CASES:
+    - Table cells should NOT have individual boxes; the entire table is one box.
+    - Page headers/footers as separate regions.
+    - Text wrapped around images as distinct regions.
 
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
+OUTPUT FORMAT:
+Return a JSON list of objects. Each object MUST have:
+- "box": A list of 4 normalized coordinates [xmin, ymin, xmax, ymax].
+- "text": The extracted text string from that box. Ensure all text within the box is captured.
 
-In the meantime, below is an example of what you can do with just a few lines of code:
+Example of a valid JSON response:
+[
+  {"box": [0.070, 0.120, 0.930, 0.280], "text": "Document Title and Header Information"},
+  {"box": [0.120, 0.350, 0.880, 0.650], "text": "Table content including headers and all rows..."},
+  {"box": [0.100, 0.700, 0.900, 0.850], "text": "This is the first paragraph of text, potentially spanning multiple lines but grouped as one logical block."},
+  {"box": [0.100, 0.880, 0.900, 0.950], "text": "Another paragraph or distinct text block."}
+]
+
+ONLY RETURN THE VALID JSON LIST. No explanations, apologies, or other text outside the JSON structure.
+If no text regions are found, return an empty JSON list: [].
 """
 
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+# --- Helper Functions ---
+
+def get_gemini_api_key() -> Optional[str]:
+    """Gets the Gemini API key from Streamlit secrets, environment variables, or user input."""
+    if "GOOGLE_API_KEY" in st.secrets:
+        return st.secrets["GOOGLE_API_KEY"]
+    api_key_env = os.getenv("GOOGLE_API_KEY")
+    if api_key_env:
+        return api_key_env
+    
+    st.sidebar.warning("Google API Key not found in secrets or environment variables.")
+    api_key_input = st.sidebar.text_input(
+        "Enter your Google API Key:", type="password", key="api_key_input"
+    )
+    if api_key_input:
+        st.session_state.GOOGLE_API_KEY = api_key_input
+        return api_key_input
+    return None
+
+def parse_gemini_response(response_text: str) -> List[Dict[str, Any]]:
+    """
+    Parses the Gemini response to extract a list of dicts with "box" and "text".
+    Tries to load as JSON, falls back to regex if needed (though ideally not).
+    """
+    try:
+        # Attempt to find JSON block if model wraps it in markdown
+        match = re.search(r"```json\s*([\s\S]*?)\s*```", response_text, re.DOTALL)
+        if match:
+            json_str = match.group(1)
+        else:
+            # Assume the response is plain JSON or a Python list string
+            json_str = response_text.strip()
+            if not (json_str.startswith('[') and json_str.endswith(']')):
+                 # If it doesn't look like a list, try to find the list within the text
+                list_match = re.search(r'(\[[\s\S]*\])', json_str)
+                if list_match:
+                    json_str = list_match.group(1)
+                else:
+                    st.warning(f"Response doesn't appear to be a list: {response_text[:200]}...")
+                    return []
+
+
+        # Try direct JSON parsing
+        parsed_data = json.loads(json_str)
+        if isinstance(parsed_data, list):
+            # Validate structure
+            validated_data = []
+            for item in parsed_data:
+                if isinstance(item, dict) and "box" in item and "text" in item and \
+                   isinstance(item["box"], list) and len(item["box"]) == 4:
+                    validated_data.append(item)
+                else:
+                    st.warning(f"Skipping invalid item in JSON: {item}")
+            return validated_data
+        else:
+            st.warning(f"Parsed JSON is not a list: {type(parsed_data)}")
+            return []
+
+    except json.JSONDecodeError as e:
+        st.warning(f"JSONDecodeError: {e}. Raw response: {response_text[:500]}")
+        # Fallback to regex if JSON parsing fails (less robust)
+        # This regex assumes the "box" and "text" structure as described in the prompt
+        # It's a simplified regex for demonstration and might need adjustment for complex texts.
+        regions = []
+        # Regex to find "box": [coords], "text": "content"
+        # This is a simplified regex and might struggle with escaped quotes in text
+        pattern = r'\{\s*"box"\s*:\s*\[([\d\.]+),\s*([\d\.]+),\s*([\d\.]+),\s*([\d\.]+)\]\s*,\s*"text"\s*:\s*"(.*?)"\s*\}'
+        matches = re.findall(pattern, response_text, re.DOTALL)
+        for match in matches:
+            try:
+                box = [float(c) for c in match[:4]]
+                text = match[4].replace('\\n', '\n').replace('\\"', '"') # Handle basic escapes
+                regions.append({"box": box, "text": text})
+            except ValueError:
+                st.warning(f"Could not parse box coordinates from regex match: {match}")
+        if not regions and response_text.strip() and response_text.strip() != "[]":
+             st.error(f"Failed to parse response using JSON and regex. Raw: {response_text[:200]}")
+        return regions
+    except Exception as e:
+        st.error(f"Error parsing Gemini response: {e}. Raw response: {response_text[:500]}")
+        return []
+
+
+def draw_bounding_boxes(image: Image.Image, regions: List[Dict[str, Any]]) -> Image.Image:
+    """Draws numbered bounding boxes on the image."""
+    if not regions:
+        return image
+
+    draw = ImageDraw.Draw(image)
+    width, height = image.size
+    
+    try:
+        # Try to load a common font, fall back to default
+        font = ImageFont.truetype("arial.ttf", int(height * 0.02))  # Adjust size as needed
+    except IOError:
+        font = ImageFont.load_default()
+
+    for i, region_data in enumerate(regions):
+        try:
+            box = region_data.get("box")
+            if not (isinstance(box, list) and len(box) == 4):
+                st.warning(f"Skipping invalid box data for region {i+1}: {box}")
+                continue
+
+            # Convert normalized coordinates to pixel values, clamping to image bounds
+            xmin = max(0.0, min(1.0, box[0])) * width
+            ymin = max(0.0, min(1.0, box[1])) * height
+            xmax = max(0.0, min(1.0, box[2])) * width
+            ymax = max(0.0, min(1.0, box[3])) * height
+
+            if xmin >= xmax or ymin >= ymax:
+                st.warning(f"Skipping invalid box dimensions for region {i+1}: {[xmin, ymin, xmax, ymax]}")
+                continue
+
+            draw.rectangle([xmin, ymin, xmax, ymax], outline="#00FF00", width=3)
+            label = str(i + 1)
+            
+            # Position label inside the box, handle potential small boxes
+            text_x = xmin + 5
+            text_y = ymin + 5
+            
+            # For very small boxes, drawing text might be an issue, but let's try
+            text_bbox = draw.textbbox((text_x, text_y), label, font=font)
+            text_width = text_bbox[2] - text_bbox[0]
+            text_height = text_bbox[3] - text_bbox[1]
+
+            # Simple check to ensure text fits somewhat
+            if text_x + text_width > xmax - 5:
+                text_x = max(xmin, xmax - text_width - 5)
+            if text_y + text_height > ymax - 5:
+                text_y = max(ymin, ymax - text_height - 5)
+            
+            draw.text((text_x, text_y), label, fill="red", font=font)
+
+        except Exception as e:
+            st.error(f"Error drawing box for region {i+1}: {str(e)}")
+    return image
+
+def process_image_with_gemini(
+    client: genai.GenerativeModel, image_bytes: bytes
+) -> Tuple[List[Dict[str, Any]], str]:
+    """Sends image to Gemini and gets bounding boxes and text."""
+    try:
+        response = client.generate_content(
+            contents=[
+                UNIFIED_DETECTION_EXTRACTION_PROMPT,
+                {"mime_type": "image/png", "data": image_bytes}
+            ],
+            generation_config=GenerationConfig(
+                temperature=0.1, # Lower temperature for more deterministic output
+                # max_output_tokens=8192 # Max for flash 1.5
+            ),
+            safety_settings={ # Adjust as needed
+                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+            }
+        )
+        if response.parts:
+            raw_text_response = response.text # .text often combines parts
+        else: # Fallback if .text is empty but candidates exist
+            raw_text_response = response.candidates[0].content.parts[0].text if response.candidates and response.candidates[0].content.parts else ""
+        
+        if not raw_text_response:
+            st.error("Received empty response from Gemini API.")
+            return [], ""
+            
+        regions = parse_gemini_response(raw_text_response)
+        return regions, raw_text_response
+    except Exception as e:
+        st.error(f"Error calling Gemini API: {str(e)}")
+        if hasattr(e, 'response') and e.response: # For google.api_core.exceptions
+            st.error(f"API Response Error Details: {e.response}")
+        return [], f"API Error: {str(e)}"
+
+# --- Streamlit UI ---
+st.set_page_config(layout="wide")
+st.title("📄 PDF Text Detection & Extraction with Gemini")
+
+# --- API Key Configuration ---
+st.sidebar.header("🔑 API Configuration")
+api_key = get_gemini_api_key()
+gemini_client = None
+
+if api_key:
+    try:
+        genai.configure(api_key=api_key)
+        gemini_client = genai.GenerativeModel(MODEL_NAME)
+        st.sidebar.success("Gemini API Key configured.")
+    except Exception as e:
+        st.sidebar.error(f"Failed to configure Gemini: {e}")
+        api_key = None # Invalidate API key if configuration fails
+else:
+    st.sidebar.info("Please provide your Google API Key to use the application.")
+    st.info("Please enter your Google API Key in the sidebar to proceed.")
+
+
+# --- File Upload and Processing ---
+st.header("📤 Upload PDF")
+uploaded_file = st.file_uploader("Choose a PDF file", type=["pdf"])
+dpi_options = [150, 200, 300, 400]
+selected_dpi = st.select_slider(
+    "Select DPI for PDF to Image conversion:",
+    options=dpi_options,
+    value=200 # Default DPI
+)
+
+if uploaded_file and st.button("🚀 Analyze PDF", disabled=not api_key or not gemini_client):
+    if not api_key or not gemini_client:
+        st.error("API Key not configured. Please enter it in the sidebar.")
+    else:
+        with st.spinner(f"Processing PDF ({uploaded_file.name})... This may take a moment."):
+            try:
+                pdf_bytes = uploaded_file.read()
+                pil_images: List[Image.Image] = convert_from_bytes(pdf_bytes, dpi=selected_dpi)
+                
+                if not pil_images:
+                    st.warning("Could not extract any images from the PDF.")
+                else:
+                    st.success(f"PDF converted to {len(pil_images)} image(s). Analyzing with Gemini...")
+
+                    page_results = []
+                    for i, p_image in enumerate(pil_images):
+                        page_progess = st.progress(0, text=f"Processing Page {i+1}/{len(pil_images)}...")
+                        
+                        # Convert PIL image to bytes for API
+                        img_byte_arr = io.BytesIO()
+                        p_image.save(img_byte_arr, format='PNG')
+                        img_bytes = img_byte_arr.getvalue()
+                        page_progess.progress(30, text=f"Page {i+1}: Sending to Gemini...")
+
+                        regions, raw_response = process_image_with_gemini(gemini_client, img_bytes)
+                        page_progess.progress(90, text=f"Page {i+1}: Received response.")
+                        
+                        page_results.append({
+                            "original_image": p_image,
+                            "regions": regions,
+                            "raw_response": raw_response
+                        })
+                        page_progess.progress(100, text=f"Page {i+1}: Done.")
+                        page_progess.empty()
+
+                    # Store results in session state to avoid re-processing on minor UI interaction
+                    # Though with the button click, this is less critical for single runs.
+                    st.session_state.page_results = page_results
+
+            except Exception as e:
+                st.error(f"An error occurred during PDF processing: {str(e)}")
+                if "page_results" in st.session_state:
+                    del st.session_state.page_results # Clear partial results
+
+# --- Display Results ---
+if "page_results" in st.session_state and st.session_state.page_results:
+    st.header("📊 Analysis Results")
+    results = st.session_state.page_results
+    
+    tab_titles = [f"Page {i+1}" for i in range(len(results))]
+    tabs = st.tabs(tab_titles)
+
+    for idx, (tab, page_data) in enumerate(zip(tabs, results)):
+        with tab:
+            st.subheader(f"Page {idx + 1} Analysis")
+            original_image = page_data["original_image"]
+            regions = page_data["regions"]
+            raw_response = page_data["raw_response"]
+
+            col1, col2 = st.columns(2)
+
+            with col1:
+                st.image(original_image, caption="Original Image", use_container_width=True)
+
+            with col2:
+                if regions:
+                    annotated_image = draw_bounding_boxes(original_image.copy(), regions)
+                    st.image(annotated_image, caption=f"Detected {len(regions)} Text Regions", use_container_width=True)
+                else:
+                    st.image(original_image.copy(), caption="No text regions detected by model", use_container_width=True)
+                    st.warning("No text regions were successfully parsed from the model's response for this page.")
+
+            if regions:
+                st.subheader("📝 Extracted Texts & Regions")
+                for i, region_data in enumerate(regions):
+                    box = region_data.get("box", "N/A")
+                    text_content = region_data.get("text", "No text extracted for this region.")
+                    with st.expander(f"Region {i+1} (Box: {box})", expanded= (i<3) ): # Expand first 3 by default
+                        st.markdown(f"**Coordinates:** `{box}`")
+                        st.markdown("**Text:**")
+                        st.text_area(f"text_area_{idx}_{i}", text_content, height=max(100, int(len(text_content)*0.5)), disabled=True, label_visibility="collapsed")
+            else:
+                st.info("No text regions to display for this page.")
+
+            with st.expander("🔍 Debug: Raw Gemini API Response", expanded=False):
+                st.code(raw_response, language='json')
+
+elif uploaded_file and not api_key:
+    st.warning("Please enter your API key in the sidebar and click 'Analyze PDF' again.")
+
+st.markdown("---")
+st.markdown("Developed with Gemini & Streamlit.")