Update page_files/Upload_Data.py

page_files/Upload_Data.py

@@ -27,10 +27,17 @@ from PIL import Image
 from dotenv import load_dotenv
 load_dotenv()
 
-_GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "")
+_GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+if not _GEMINI_API_KEY:
+    raise RuntimeError("GEMINI_API_KEY not set in environment")
 
 # ── imports from doctodb_rag (data extraction) ────────────────────────────────
-from categorized.Backend.PDF_DataExtraction import run_pipeline
+from categorized.Backend.FinalVerdict import (
+    call_gemini_from_bytes,
+    convert_to_dataframe,
+    _extract_sentences,
+    verify_dataframe,
+)
 
 # ── imports from figure_extractor (image extraction) ─────────────────────────
 from categorized.Backend.Pdf_ImageExtraction import (
@@ -161,17 +168,18 @@ def save_single_image_with_property(
 # expected by the rest of the UI (list of {caption, page, image_data}).
 # ─────────────────────────────────────────────────────────────────────────────
 
+_GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "AIzaSyBzyMFKEqcjsWpR-OGAY42T250o1O39v3Y")
 
 def extract_images(pdf_path: str) -> list:
+    """
+    Use figure_extractor to detect and crop plot images from a PDF path.
+    Returns a list compatible with the image_results shape used throughout the UI:
+      [{ "caption": str, "page": int, "image_data": [{"array": bgr_ndarray, "filename": str}] }]
+    """
     try:
-        from categorized.Backend.Pdf_ImageExtraction import get_available_model
-        import google.generativeai as genai
-        api_key = os.getenv("GEMINI_API_KEY", "")
-        genai.configure(api_key=api_key)
-        model_name = get_available_model(api_key)
-        active_model = genai.GenerativeModel(model_name)
-        plot_data = get_plot_data_from_llm(active_model, pdf_path)
-        raw_plots = extract_plots(
+        # gemini_model = init_gemini(_GEMINI_API_KEY)
+        plot_data    = get_plot_data_from_llm( GEMINI_MODEL, pdf_path)
+        raw_plots    = extract_plots(
             pdf_path=pdf_path,
             plot_data=plot_data,
             pad=22,
@@ -181,23 +189,32 @@ def extract_images(pdf_path: str) -> list:
         log.error(f"extract_images failed: {e}")
         return []
 
+
+
+
+    # raw_plots items: {caption, page, path, plot_score, plot_type}
+    # Convert to image_results shape
     image_results = []
     for item in raw_plots:
         bgr = cv2.imread(item["path"]) if item.get("path") else None
+        # clean up temp file written by extract_plots
         if item.get("path") and os.path.exists(item["path"]):
             try:
                 os.remove(item["path"])
             except Exception:
                 pass
+
         page    = item.get("page", 1)
         caption = item.get("caption", f"Figure (page {page})")
         safe    = re.sub(r"[^\w\-]", "_", caption)[:40]
         filename = f"page{page}_{safe}.png"
+
         image_results.append({
             "caption":    caption,
             "page":       page,
             "image_data": [{"array": bgr, "filename": filename}] if bgr is not None else [],
         })
+
     return image_results
 
 
@@ -688,9 +705,22 @@ def render_material_data_tab(pdf_path: str):
         with open(pdf_path, "rb") as f:
             pdf_bytes = f.read()
 
-        df, _, _, _, api_errors, meta = run_pipeline(
-            pdf_bytes, progress_callback=_cb
-        )
+        _cb("Extracting via Gemini…", 0.30)
+        data = call_gemini_from_bytes(pdf_bytes)
+        df = convert_to_dataframe(data)
+        api_errors = []
+
+        if not df.empty:
+            _cb("Verifying against source PDF…", 0.70)
+            sentences = _extract_sentences(pdf_bytes)
+            df = verify_dataframe(df, sentences)
+
+        row0 = df.iloc[0] if not df.empty else {}
+        meta = {
+            "material_name": str(row0.get("material_name", "")) if not df.empty else "",
+            "material_abbreviation": str(row0.get("material_abbreviation", "")) if not df.empty else "",
+        }
+        _cb("Done.", 1.0)
         elapsed_total = time.time() - start_ts
         bar.progress(1.0)
         status.empty()