Upload Pdf_ImageExtraction.py

page_files/categorized/Backend/Pdf_ImageExtraction.py

@@ -1,390 +1,379 @@
-import os
-import re
+import logging
+
+import requests
+
+log = logging.getLogger(__name__)
+
+import streamlit as st
+import google.generativeai as genai
+import fitz  # PyMuPDF
 import json
-import math
 import tempfile
-import fitz  # PyMuPDF
-import cv2
+import re
+import os
+
 import numpy as np
-from PIL import Image
-import streamlit as st
+import cv2
+from typing import Optional, Tuple
 
-# -------------------
-# Config
-# -------------------
-DPI = 300
-OUT_DIR = "outputs"
 
-KEEP_ONLY_STRESS_STRAIN = False
 
-CAP_RE = re.compile(r"^(Fig\.?\s*\d+|Figure\s*\d+)\b", re.IGNORECASE)
-SS_KW  = re.compile(
-    r"(stress\s*[-–]?\s*strain|stress|strain|tensile|MPa|GPa|kN|yield|elongation)",
-    re.IGNORECASE
-)
+# ----------------------------
+# Gemini init (more deterministic)
+# ----------------------------
+API_KEY = os.getenv("GEMINI_API_KEY", " ")
 
-# -------------------
-# Render helpers
-# -------------------
-def render_page(page, dpi=DPI):
-    mat = fitz.Matrix(dpi/72, dpi/72)
-    pix = page.get_pixmap(matrix=mat, alpha=False)
-    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-    return img, mat
-
-def pdf_to_px_bbox(bbox_pdf, mat):
-    x0, y0, x1, y1 = bbox_pdf
-    sx, sy = mat.a, mat.d
-    return (int(float(x0) * sx), int(float(y0) * sy), int(float(x1) * sx), int(float(y1) * sy))
-
-def safe_crop_px(pil_img, box):
-    if not isinstance(box, (tuple, list)):
-        return None
-    if len(box) == 1 and isinstance(box[0], (tuple, list)) and len(box[0]) == 4:
-        box = box[0]
-    if len(box) != 4:
-        return None
-
-    x0, y0, x1, y1 = box
-    if any(isinstance(v, (tuple, list)) for v in (x0, y0, x1, y1)):
-        return None
+_cached_model: Optional[str] = None
+ 
 
+def get_available_model(api_key: str) -> str:
+    """Fetch available models and return best flash model found."""
+    global _cached_model
+    if _cached_model:
+        return _cached_model
     try:
-        x0 = int(x0)
-        y0 = int(y0)
-        x1 = int(x1)
-        y1 = int(y1)
-    except (TypeError, ValueError):
-        return None
-
-    if x1 < x0: 
-        x0, x1 = x1, x0
-    if y1 < y0: 
-        y0, y1 = y1, y0
-
-    W, H = pil_img.size
-    x0 = max(0, min(W, x0))
-    x1 = max(0, min(W, x1))
-    y0 = max(0, min(H, y0))
-    y1 = max(0, min(H, y1))
-    if x1 <= x0 or y1 <= y0:
-        return None
-    return pil_img.crop((x0, y0, x1, y1))
-
-# -------------------
-# Captions
-# -------------------
-def find_caption_blocks(page):
-    caps = []
-    blocks = page.get_text("blocks")
-    for b in blocks:
-        x0, y0, x1, y1, text = b[0], b[1], b[2], b[3], b[4]
-        t = " ".join(str(text).strip().split())
-        if CAP_RE.match(t):
-            caps.append({"bbox": (x0, y0, x1, y1), "text": t})
-    return caps
-
-# -------------------
-# Dedupe: dHash
-# -------------------
-def dhash64(pil_img):
-    gray = pil_img.convert("L").resize((9, 8), Image.LANCZOS)
-    pixels = list(gray.getdata())
-    bits = 0
-    for r in range(8):
-        for c in range(8):
-            left = pixels[r * 9 + c]
-            right = pixels[r * 9 + c + 1]
-            bits = (bits << 1) | (1 if left > right else 0)
-    return bits
-
-# -------------------
-# Rejectors
-# -------------------
-def has_colorbar_like_strip(pil_img):
-    img = np.array(pil_img)
-    if img.ndim != 3:
-        return False
-    H, W, _ = img.shape
-    if W < 250 or H < 150:
-        return False
-    strip_w = max(18, int(0.07 * W))
-    strip = img[:, W-strip_w:W, :]
-    q = (strip // 24).reshape(-1, 3)
-    uniq = np.unique(q, axis=0)
-    return len(uniq) > 70
-
-def texture_score(pil_img):
-    gray = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2GRAY)
-    lap = cv2.Laplacian(gray, cv2.CV_64F)
-    return float(lap.var())
-
-def is_mostly_legend(pil_img):
-    gray = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2GRAY)
-    bw = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
-    bw = cv2.medianBlur(bw, 3)
-    H, W = bw.shape
-    fill = float(np.count_nonzero(bw)) / float(H * W)
-    return (0.03 < fill < 0.18) and (min(H, W) < 260)
-
-# -------------------
-# Plot detection
-# -------------------
-def detect_axes_lines(pil_img):
-    gray = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2GRAY)
-    edges = cv2.Canny(gray, 50, 150)
-    H, W = gray.shape
-    min_len = int(0.28 * min(H, W))
-
-    lines = cv2.HoughLinesP(
-        edges, 1, np.pi/180,
-        threshold=90,
-        minLineLength=min_len,
-        maxLineGap=14
+        resp = requests.get(
+            f"https://generativelanguage.googleapis.com/v1beta/models?key={api_key}",
+            timeout=10
+        )
+        if not resp.ok:
+            _cached_model = "gemini-2.0-flash"
+            return _cached_model
+        models = resp.json().get("models", [])
+        model_names = [m["name"].replace("models/", "") for m in models
+                      if "generateContent" in m.get("supportedGenerationMethods", [])]
+        preferred = [
+            "gemini-2.5-flash",
+            "gemini-2.5-flash-preview",
+            "gemini-2.0-flash",
+            "gemini-2.0-flash-lite",
+        ]
+        for pref in preferred:
+            for name in model_names:
+                if name.startswith(pref):
+                    log.info(f"Auto-selected model: {name}")
+                    _cached_model = name
+                    return _cached_model
+        flash_models = [n for n in model_names if "flash" in n]
+        _cached_model = flash_models[0] if flash_models else "gemini-2.0-flash"
+        return _cached_model
+    except Exception as e:
+        log.warning(f"Model auto-detect failed: {e}")
+        _cached_model = "gemini-2.0-flash"
+        return _cached_model
+ 
+GEMINI_MODEL = get_available_model(API_KEY)
+API_URL = (
+    f"https://generativelanguage.googleapis.com/v1beta/"
+    f"models/{GEMINI_MODEL}:generateContent"
+    f"?key={API_KEY}"
+)
+
+
+
+
+def robust_json_load(text: str):
+    """
+    Handles cases where Gemini returns JSON + extra text.
+    """
+    if not text:
+        raise ValueError("Empty model output")
+
+    cleaned = text.strip().replace("```json", "```").replace("```", "").strip()
+
+    # Find first JSON start
+    s1 = cleaned.find("[")
+    s2 = cleaned.find("{")
+    if s1 == -1 and s2 == -1:
+        raise ValueError("No JSON found in model output")
+
+    start = s1 if (s1 != -1 and (s2 == -1 or s1 < s2)) else s2
+    dec = json.JSONDecoder()
+    parsed, _ = dec.raw_decode(cleaned[start:])
+    return parsed
+
+
+def get_plot_data_from_llm(GEMINI_MODEL, pdf_path: str):
+    sample_file = genai.upload_file(path=pdf_path)
+
+    prompt = prompt = """
+Analyze this PDF and identify ONLY true data plots:
+(Line, Scatter, Bar, Histogram, Heatmap, Box, Violin).
+EXCLUDE: flowcharts, block diagrams, photos, icons, gauges/dials, 
+graphical abstracts, tables, and pure text regions.
+
+DETECTION RULES:
+1) INCLUDE ALL SUBPLOTS: If a figure has panels like (a), (b), (c), 
+   return ONE bounding box covering the ENTIRE figure group including 
+   all panels, axes, and legends.
+2) CROP AREA: box_2d MUST include:
+   - The full plot area (all data points, bars, lines)
+   - ALL axis lines and tick marks
+   - ALL axis labels and tick labels  
+   - The legend if present
+   - A small margin above the top of the plot
+   box_2d MUST EXCLUDE figure caption text below the figure.
+3) IMPORTANT: Make the bounding box GENEROUS — it is better to include 
+   too much whitespace than to cut off any part of the plot.
+4) DO NOT return tables or paragraphs as images.
+5) Skip graphical abstracts and decorative figures.
+6) Extract the FULL caption text into "caption".
+
+COORDINATES:
+- box_2d is [ymin, xmin, ymax, xmax] in 0..1000 normalized page coordinates.
+- ymin should be ABOVE the top of the plot area.
+- ymax should be BELOW the bottom axis but ABOVE the caption text.
+- xmin should be LEFT of the y-axis labels.
+- xmax should be RIGHT of the rightmost data point or legend.
+
+Return ONLY a raw JSON array like:
+[
+  {
+    "caption":"...", 
+    "page": 1, 
+    "box_2d":[ymin,xmin,ymax,xmax], 
+    "figure_kind":"plot|nonplot", 
+    "plot_type":"scatter|line|bar|hist|heatmap|box|violin|other"
+  }
+]
+IMPORTANT: Only include items where figure_kind="plot".
+"""
+
+    # Try to force JSON + reduce randomness
+    generation_config = {
+        "temperature": 0.0,
+        "top_p": 1.0,
+        "top_k": 1,
+        "max_output_tokens": 4096,
+    }
+
+    # Some environments accept response_mime_type; if yours errors, remove it.
+    response = GEMINI_MODEL.generate_content(
+        [sample_file, prompt],
+        generation_config=generation_config
     )
-    if lines is None:
-        return None, None
-
-    horizontals, verticals = [], []
-    for x1, y1, x2, y2 in lines[:, 0]:
-        dx, dy = abs(x2-x1), abs(y2-y1)
-        length = math.hypot(dx, dy)
-        if dy < 18 and dx > 0.35 * W:
-            horizontals.append((length, (x1, y1, x2, y2)))
-        if dx < 18 and dy > 0.35 * H:
-            verticals.append((length, (x1, y1, x2, y2)))
-
-    if not horizontals or not verticals:
-        return None, None
-
-    horizontals.sort(key=lambda t: t[0], reverse=True)
-    verticals.sort(key=lambda t: t[0], reverse=True)
-    return horizontals[0][1], verticals[0][1]
-
-def axis_intersection_ok(x_axis, y_axis, W, H):
-    xa_y = int(round((x_axis[1] + x_axis[3]) / 2))
-    ya_x = int(round((y_axis[0] + y_axis[2]) / 2))
-    if not (0 <= xa_y < H and 0 <= ya_x < W):
-        return False
-    if ya_x > int(0.95 * W) or xa_y < int(0.05 * H):
-        return False
-    return True
-
-def tick_text_presence_score(pil_img, x_axis, y_axis):
-    img = np.array(pil_img)
-    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-    bw = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)[1]
-    bw = cv2.medianBlur(bw, 3)
-
-    H, W = gray.shape
-    xa_y = int(round((x_axis[1] + x_axis[3]) / 2))
-    ya_x = int(round((y_axis[0] + y_axis[2]) / 2))
-
-    y0a = max(0, xa_y - 40)
-    y1a = min(H, xa_y + 110)
-    x_roi = bw[y0a:y1a, 0:W]
-
-    x0b = max(0, ya_x - 180)
-    x1b = min(W, ya_x + 50)
-    y_roi = bw[0:H, x0b:x1b]
-
-    def count_small_components(mask):
-        num, _, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
-        cnt = 0
-        for i in range(1, num):
-            x, y, w, h, area = stats[i]
-            if 4 <= w <= 150 and 4 <= h <= 150 and 20 <= area <= 5000:
-                cnt += 1
-        return cnt
-
-    return count_small_components(x_roi) + count_small_components(y_roi)
-
-def is_real_plot(pil_img):
-    if has_colorbar_like_strip(pil_img):
-        return False
-    if is_mostly_legend(pil_img):
-        return False
-
-    x_axis, y_axis = detect_axes_lines(pil_img)
-    if x_axis is None or y_axis is None:
-        return False
-
-    arr = np.array(pil_img)
-    H, W = arr.shape[0], arr.shape[1]
-    if not axis_intersection_ok(x_axis, y_axis, W, H):
-        return False
-
-    if texture_score(pil_img) > 2200:
-        return False
-
-    score = tick_text_presence_score(pil_img, x_axis, y_axis)
-    return score >= 18
-
-# -------------------
-# Candidate boxes in a region
-# -------------------
-def connected_components_boxes(pil_img):
-    img_bgr = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
-    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
-    mask = (gray < 245).astype(np.uint8) * 255
-    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((7, 7), np.uint8), iterations=2)
-    num, _, stats, _ = cv2.connectedComponentsWithStats(mask, connectivity=8)
-
-    boxes = []
-    for i in range(1, num):
-        x, y, w, h, area = stats[i]
-        boxes.append((int(area), (int(x), int(y), int(x + w), int(y + h))))
-    boxes.sort(key=lambda t: t[0], reverse=True)
-    return boxes
-
-def expand_box(box, W, H, left=0.10, right=0.06, top=0.06, bottom=0.18):
-    x0, y0, x1, y1 = box
-    bw = x1 - x0
-    bh = y1 - y0
-    ex0 = max(0, int(x0 - left * bw))
-    ex1 = min(W, int(x1 + right * bw))
-    ey0 = max(0, int(y0 - top * bh))
-    ey1 = min(H, int(y1 + bottom * bh))
-    return (ex0, ey0, ex1, ey1)
-
-# -------------------
-# Crop plot from caption
-# -------------------
-def crop_plot_from_caption(page_img, cap_bbox_pdf, mat):
-    cap_px = pdf_to_px_bbox(cap_bbox_pdf, mat)
-    cap_y0 = cap_px[1]
-    cap_y1 = cap_px[3]
-
-    W, H = page_img.size
-    search_top = max(0, cap_y0 - int(0.95 * H))
-    search_bot = min(H, cap_y1 + int(0.20 * H))
-    region = safe_crop_px(page_img, (0, search_top, W, search_bot))
-    if region is None:
-        return None
-
-    comps = connected_components_boxes(region)
-    best = None
-    best_area = -1
-
-    for area, box in comps[:35]:
-        x0, y0, x1, y1 = box
-        bw = x1 - x0
-        bh = y1 - y0
-        if bw < 220 or bh < 180:
-            continue
 
-        exp = expand_box(box, region.size[0], region.size[1])
-        cand = safe_crop_px(region, exp)
-        if cand is None:
-            continue
+    data = robust_json_load(response.text)
 
-        if not is_real_plot(cand):
+    # Hard filter if model still returns nonplots
+    filtered = []
+    for it in data if isinstance(data, list) else []:
+        if not isinstance(it, dict):
             continue
+        if (it.get("figure_kind") or "").strip().lower() == "plot":
+            filtered.append(it)
+    return filtered
 
-        if area > best_area:
-            best_area = area
-            best = cand
 
-    return best
+# ----------------------------
+# Caption trimming using PDF text blocks
+# ----------------------------
+CAP_RE = re.compile(r"^(fig\.?\s*\d+|figure\s*\d+)\b", re.IGNORECASE)
 
-# -------------------
-# Streamlit UI
-# -------------------
-def run_extraction(pdf_path, paper_id="uploaded_paper"):
-    out_paper = os.path.join(OUT_DIR, paper_id)
-    out_imgs = os.path.join(out_paper, "plots_with_axes")
-    os.makedirs(out_imgs, exist_ok=True)
 
-    doc = fitz.open(pdf_path)
-    results = []
-    seen = set()
-    saved = 0
+def trim_caption_region(page: fitz.Page, crop_rect: fitz.Rect) -> fitz.Rect:
+    """
+    If crop includes caption text (Fig./Figure...) near the bottom,
+    shrink ymax to exclude it.
+    """
+    blocks = page.get_text("blocks")  # (x0,y0,x1,y1,text,block_no,block_type)
+    if not blocks:
+        return crop_rect
 
-    for p in range(len(doc)):
-        page = doc[p]
-        caps = find_caption_blocks(page)
-        if not caps:
+    # Look for caption blocks that intersect crop and are below plot area
+    candidate_tops = []
+    for b in blocks:
+        x0, y0, x1, y1, text = b[0], b[1], b[2], b[3], b[4]
+        if not text:
+            continue
+        t = text.strip().replace("\n", " ")
+        if not CAP_RE.search(t):
             continue
 
-        page_img, mat = render_page(page, dpi=DPI)
-
-        for cap in caps:
-            cap_text = cap["text"]
-
-            if KEEP_ONLY_STRESS_STRAIN and not SS_KW.search(cap_text):
-                continue
+        block_rect = fitz.Rect(x0, y0, x1, y1)
+        if block_rect.intersects(crop_rect):
+            # likely caption inside crop, use its top boundary
+            candidate_tops.append(y0)
+
+    if not candidate_tops:
+        return crop_rect
+
+    cap_top = min(candidate_tops)
+    # Reduce crop bottom just above caption top
+    new = fitz.Rect(crop_rect.x0, crop_rect.y0, crop_rect.x1, min(crop_rect.y1, cap_top - 2))
+    return new if new.y1 > new.y0 + 10 else crop_rect
+
+
+# ----------------------------
+# Plot-likeness filter (reject text/diagrams)
+# ----------------------------
+def plot_likeness_score(bgr: np.ndarray) -> float:
+    """
+    Heuristic: plots usually have strong horizontal/vertical lines (axes, grids),
+    and moderate edge density. Text blocks often have many tiny components and no axes.
+    """
+    gray = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY)
+    h, w = gray.shape[:2]
+    if h < 40 or w < 40:
+        return 0.0
+
+    # edges
+    edges = cv2.Canny(gray, 60, 160)
+    edge_density = float(np.mean(edges > 0))
+
+    # detect long straight lines (axes/grid)
+    lines = cv2.HoughLinesP(edges, 1, np.pi / 180, threshold=80,
+                            minLineLength=int(min(h, w) * 0.35), maxLineGap=10)
+
+   
+    hv_lines = 0
+    if lines is not None:
+        for x1, y1, x2, y2 in lines[:, 0]:
+            dx = abs(x2 - x1)
+            dy = abs(y2 - y1)
+            if dx > 3 * dy or dy > 3 * dx:
+                hv_lines += 1
+
+    # connected components for "textiness"
+    _, bw = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
+    num_labels, _, stats, _ = cv2.connectedComponentsWithStats(bw, connectivity=8)
+    # count small-ish components (typical for dense text)
+    small = 0
+    for i in range(1, num_labels):
+        area = stats[i, cv2.CC_STAT_AREA]
+        if 10 <= area <= 120:
+            small += 1
+    small_density = small / (h * w / 10000.0 + 1e-6)
+
+    # Score: prefer hv_lines + moderate edges; penalize extreme textiness
+    score = (0.55 * min(1.0, hv_lines / 8.0)) + (0.35 * min(1.0, edge_density / 0.12)) - (0.20 * min(1.0, small_density / 25.0))
+    return float(max(0.0, min(1.0, score)))
+
+
+def looks_like_plot(bgr: np.ndarray, threshold: float = 0.35) -> Tuple[bool, float]:
+    s = plot_likeness_score(bgr)
+    return (s >= threshold), s
+
+
+# ----------------------------
+# Extraction with better padding
+# ----------------------------
+def extract_plots(pdf_path: str, plot_data: list, pad: int, score_thresh: float):
+    doc = fitz.open(pdf_path)
+    results = []
 
-            fig = crop_plot_from_caption(page_img, cap["bbox"], mat)
-            if fig is None:
-                continue
+    for idx, item in enumerate(plot_data):
+        page_num = int(item.get("page", 1)) - 1
+        if page_num < 0 or page_num >= len(doc):
+            continue
 
-            if fig.size[0] > 8 and fig.size[1] > 8:
-                fig = fig.crop((2, 2, fig.size[0]-2, fig.size[1]-2))
+        page = doc[page_num]
+        rect = page.rect
+        box = item.get("box_2d", None)
+        if not box or len(box) != 4:
+            continue
 
+        # normalized 0..1000 → absolute page coords
+        ymin = (box[0] * rect.height / 1000.0)
+        xmin = (box[1] * rect.width / 1000.0)
+        ymax = (box[2] * rect.height / 1000.0)
+        xmax = (box[3] * rect.width / 1000.0)
+
+        # Asymmetric padding: keep labels, avoid caption blow-up
+        pad_left = pad * 1.4
+        pad_right = pad * 0.9
+        pad_top = pad * 0.9
+        pad_bottom = pad * 0.7  # smaller bottom pad to avoid captions
+
+        fitz_rect = fitz.Rect(
+            max(0, xmin - pad_left),
+            max(0, ymin - pad_top),
+            min(rect.width, xmax + pad_right),
+            min(rect.height, ymax + pad_bottom)
+        )
+
+        # Trim caption if accidentally included
+        fitz_rect = trim_caption_region(page, fitz_rect)
+
+        # Render
+        pix = page.get_pixmap(clip=fitz_rect, dpi=300)
+        img_path = f"plot_res_{idx}.png"
+        pix.save(img_path)
+
+        # Post-filter: reject non-plots
+        bgr = cv2.imread(img_path)
+        ok, score = looks_like_plot(bgr, threshold=score_thresh)
+
+        if not ok:
+            # remove saved non-plot image
             try:
-                h = dhash64(fig)
+                os.remove(img_path)
             except Exception:
-                continue
-
-            if h in seen:
-                continue
-            seen.add(h)
-
-            img_name = f"p{p+1:02d}_{saved:04d}.png"
-            img_path = os.path.join(out_imgs, img_name)
-            fig.save(img_path)
+                pass
+            continue
 
-            results.append({
-                "page": p + 1,
-                "caption": cap_text,
-                "image": img_path
-            })
-            saved += 1
+        results.append({
+            "caption": item.get("caption", f"Figure {idx}"),
+            "page": item.get("page", 1),
+            "path": img_path,
+            "plot_score": round(score, 3),
+            "plot_type": item.get("plot_type", "unknown")
+        })
 
-    out_json = os.path.join(out_paper, "plots_with_axes.json")
-    with open(out_json, "w", encoding="utf-8") as f:
-        json.dump(results, f, indent=2, ensure_ascii=False)
+    doc.close()
+    return results
 
-    return results, out_json
 
+# ----------------------------
+# Streamlit UI
+# ----------------------------
 def main():
-    st.set_page_config(page_title="Research Paper Plot Extractor", layout="wide")
-    st.title(" Plot Extractor (Upload PDF)")
-
-    uploaded = st.file_uploader("Upload a research paper PDF", type=["pdf"])
-    if not uploaded:
-        st.info("Upload a PDF to extract plots.")
-        return
-
-    paper_id = os.path.splitext(uploaded.name)[0].replace(" ", "_")
-
-    with tempfile.TemporaryDirectory() as tmpdir:
-        pdf_path = os.path.join(tmpdir, uploaded.name)
-        with open(pdf_path, "wb") as f:
-            f.write(uploaded.read())
-
-        with st.spinner("Extracting plots..."):
-            results, out_json = run_extraction(pdf_path, paper_id=paper_id)
-
-        st.success(f"Extracted {len(results)} plots.")
-
-        # Show images + captions
-        for r in results:
-            st.markdown(f"**Page {r['page']}** — {r['caption']}")
-            st.image(r["image"], use_container_width=True)
-            st.divider()
-
-        # JSON viewer + download
-        st.subheader("JSON Output")
-        st.json(results)
+    st.set_page_config(layout="wide", page_title="Scientific Figure Extractor")
+    st.title("Scientific Figure & Subplot Extractor")
+
+    if "raw_data" not in st.session_state:
+        st.session_state.raw_data = None
+    if "temp_pdf" not in st.session_state:
+        st.session_state.temp_pdf = None
+
+    with st.sidebar:
+        api_key = st.text_input("Gemini API Key", type="password")
+        st.divider()
+        st.subheader("Crop & Filter")
+        crop_pad = st.slider("Padding (px-like on page coords)", 0, 80, 22)
+        score_thresh = st.slider("Reject non-plots (higher = stricter)", 0.10, 0.80, 0.20, 0.01)
+
+    uploaded_file = st.file_uploader("Upload PDF", type="pdf")
+
+    if uploaded_file and api_key:
+        if st.button("Extract All Data Figures"):
+            with st.spinner("AI is scanning for plots and sub-panels..."):
+                with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
+                    tmp.write(uploaded_file.getbuffer())
+                    st.session_state.temp_pdf = tmp.name
+
+                # model = init_gemini(api_key)
+                st.session_state.raw_data = get_plot_data_from_llm(GEMINI_MODEL, st.session_state.temp_pdf)
+                st.success(f"Detected {len(st.session_state.raw_data)} candidate plot figures (LLM).")
+
+    # Results Display
+    if st.session_state.raw_data and st.session_state.temp_pdf:
+        st.divider()
+        processed_figures = extract_plots(
+            st.session_state.temp_pdf,
+            st.session_state.raw_data,
+            pad=crop_pad,
+            score_thresh=score_thresh
+        )
+
+        st.success(f"Kept {len(processed_figures)} figures after non-plot filtering.")
+
+        for fig in processed_figures:
+            with st.container(border=True):
+                st.markdown(f"**{fig['caption']}**")
+                st.caption(f"Page {fig['page']} | plot_score={fig['plot_score']} | plot_type={fig['plot_type']}")
+                st.image(fig["path"], use_container_width=True)
 
-        with open(out_json, "rb") as f:
-            st.download_button(
-                "Download JSON",
-                data=f,
-                file_name=os.path.basename(out_json),
-                mime="application/json"
-            )
 
 if __name__ == "__main__":
-    main()
+    main()