from __future__ import annotations import os, io, re, json, time, mimetypes, tempfile, string from typing import List, Union, Tuple, Any, Iterable from PIL import Image import pandas as pd import gradio as gr import google.generativeai as genai import requests import pdfplumber # ================== CONFIG ================== DEFAULT_API_KEY = "AIzaSyBbK-1P3JD6HPyE3QLhkOps6_-Xo3wUFbs" INTERNAL_MODEL_MAP = { "Gemini 2.5 Flash": "gemini-2.5-flash", "Gemini 2.5 Pro": "gemini-2.5-pro", } EXTERNAL_MODEL_NAME = "prithivMLmods/Camel-Doc-OCR-062825 (External)" try: RESAMPLE = Image.Resampling.LANCZOS except AttributeError: RESAMPLE = Image.LANCZOS PROMPT_FREIGHT_JSON = """ Please analyze the freight rate table in the file I provide and convert it into JSON in the following structure: { "shipping_line": "...", "shipping_line_code": "...", "shipping_line_reason": "Why this carrier is chosen?", "fee_type": "Air Freight", "valid_from": ..., "valid_to": ..., "charges": [ { "frequency": "...", "package_type": "...", "aircraft_type": "...", "direction": "Export or Import or null", "origin": "...", "destination": "...", "charge_name": "...", "charge_code": "...", "charge_code_reason": "...", "cargo_type": "...", "currency": "...", "transit": "...", "transit_time": "...", "weight_breaks": { "M": ..., "N": ..., "+45kg": ..., "+100kg": ..., "+300kg": ..., "+500kg": ..., "+1000kg": ..., "other": { key: value }, "weight_breaks_reason":"Why chosen weight_breaks?" }, "remark": "..." } ], "local_charges": [ { "charge_name": "...", "charge_code": "...", "unit": "...", "amount": ..., "remark": "..." } ] } ### Date rules - valid_from format: - `DD/MM/YYYY` (if full date) - `01/MM/YYYY` (if month+year only) - `01/01/YYYY` (if year only) - `UFN` if missing - valid_to: - exact `DD/MM/YYYY` if present - else `UFN` STRICT RULES: - ONLY return a single JSON object as specified above. - All rates must exactly match the corresponding weight break columns (M,N,45kg, 100kg, 300kg, 500kg, 1000kg, etc.). set null if N/A. No assumptions or interpolations. - If the table shows "RQ" or similar, set value as "RQST". - Group same-price destinations into one record separated by "/". - Always use IATA code for origin and destination. - Flight number (e.g. ZH118) is not charge code. - Frequency: D[1-7]; 'Daily' = D1234567. Join multiple (e.g. D3,D4→D34). - If local charges exist, list them. - If validity missing, set null. - Direction: Export if origin is Vietnam (SGN, HAN, DAD...), else Import. - Provide short plain English reasons for "shipping_line_reason" & "charge_code_reason". - Replace commas in remarks with semicolons. - Only return JSON. """ # ================== HELPERS ================== import fitz # PyMuPDF def pdf_to_images(pdf_bytes: bytes) -> list[Image.Image]: doc = fitz.open(stream=pdf_bytes, filetype="pdf") pages = [] for p in doc: pix = p.get_pixmap(dpi=200) img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) pages.append(img) return pages def ensure_rgb(im: Image.Image) -> Image.Image: return im.convert("RGB") if im.mode != "RGB" else im def _read_file_bytes(upload: Union[str, os.PathLike, dict, object] | None) -> bytes: if upload is None: raise ValueError("No file uploaded.") if isinstance(upload, (str, os.PathLike)): with open(upload, "rb") as f: return f.read() if isinstance(upload, dict) and "path" in upload: with open(upload["path"], "rb") as f: return f.read() if hasattr(upload, "read"): return upload.read() raise TypeError(f"Unsupported file object: {type(upload)}") def _guess_name_and_mime(file, file_bytes: bytes) -> Tuple[str, str]: if isinstance(file, (str, os.PathLike)): filename = os.path.basename(str(file)) elif isinstance(file, dict) and "name" in file: filename = os.path.basename(file["name"]) elif isinstance(file, dict) and "path" in file: filename = os.path.basename(file["path"]) else: filename = "upload.bin" mime, _ = mimetypes.guess_type(filename) if not mime: if len(file_bytes) >= 4 and file_bytes[:4] == b"%PDF": mime = "application/pdf" if not filename.lower().endswith(".pdf"): filename += ".pdf" else: mime = "image/png" return filename, mime # ================== PDF CHECK STEP ================== def check_pdf_structure(file_bytes: bytes) -> str: """Kiểm tra nhanh file PDF có phải bảng nhiều cột, nhiều trang không.""" try: with pdfplumber.open(io.BytesIO(file_bytes)) as pdf: if len(pdf.pages) <= 2: return "không" table_pages = 0 for page in pdf.pages[:3]: tables = page.find_tables() if tables and len(tables) > 0: table_pages += 1 if table_pages >= 1: return "có" text = "\n".join([(p.extract_text() or "") for p in pdf.pages[:2]]) num_tokens = sum(ch.isdigit() for ch in text) line_count = len(text.splitlines()) if num_tokens > 100 and line_count > 20: return "có" return "không" except Exception as e: print("PDF check error:", e) return "không" # ================== OCR CORE (Gemini) ================== def run_process_internal_base_v2(file_bytes, filename, mime, question, model_choice, temperature, top_p, batch_size=3): api_key = os.environ.get("GOOGLE_API_KEY", DEFAULT_API_KEY) if not api_key: return "ERROR: Missing GOOGLE_API_KEY.", None genai.configure(api_key=api_key) model_name = INTERNAL_MODEL_MAP.get(model_choice, "gemini-2.5-flash") model = genai.GenerativeModel(model_name=model_name, generation_config={"temperature": float(temperature), "top_p": float(top_p)}) if file_bytes[:4] == b"%PDF": pages = pdf_to_images(file_bytes) else: pages = [Image.open(io.BytesIO(file_bytes))] user_prompt = (question or "").strip() or PROMPT_FREIGHT_JSON all_json_results, all_text_results = [], [] previous_header_json = None def _safe_text(resp): try: return resp.text except: return "" for i in range(0, len(pages), batch_size): batch = pages[i:i+batch_size] uploaded = [] for im in batch: with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp: im.save(tmp.name) up = genai.upload_file(path=tmp.name, mime_type="image/png") up = genai.get_file(up.name) uploaded.append(up) context_prompt = user_prompt resp = model.generate_content([context_prompt] + uploaded) text = _safe_text(resp) all_text_results.append(text) for up in uploaded: try: genai.delete_file(up.name) except: pass return "\n\n".join(all_text_results), None # ================== EXTERNAL API (nếu có) ================== def run_process_external(file_bytes, filename, mime, question, api_url, temperature, top_p): if not api_url: return "ERROR: Missing external API endpoint.", None data = {"prompt": question or "", "temperature": str(temperature), "top_p": str(top_p)} files = {"file": (filename, file_bytes, mime)} r = requests.post(api_url, files=files, data=data, timeout=60) if r.status_code >= 400: return f"ERROR: External API HTTP {r.status_code}: {r.text[:200]}", None return r.text, None # ================== MAIN ROUTER (đã thêm STEP CHECK) ================== def run_process(file, question, model_choice, temperature, top_p, external_api_url): """ Router (có bước kiểm tra PDF/table trước khi xử lý): - Nếu PDF nhiều trang/nhiều bảng -> extract trước (pdfplumber) - Ngược lại -> OCR trực tiếp Gemini """ try: if file is None: return "ERROR: No file uploaded.", None file_bytes = _read_file_bytes(file) filename, mime = _guess_name_and_mime(file, file_bytes) # STEP 1️⃣: Check PDF structure if mime == "application/pdf" or file_bytes[:4] == b"%PDF": check_result = check_pdf_structure(file_bytes) print(f"[PDF Check] {filename}: {check_result}") if check_result == "có": try: print("➡️ PDF có nhiều cột/nhiều trang → dùng pdfplumber extract trước rồi Gemini.") all_dfs = [] saved_header = None with pdfplumber.open(io.BytesIO(file_bytes)) as pdf: for page_idx, page in enumerate(pdf.pages, start=1): print(f"📄 Đang xử lý trang {page_idx}...") table = page.extract_table({ "vertical_strategy": "lines", "horizontal_strategy": "text", "snap_tolerance": 3, "intersection_tolerance": 5, }) if not table or len(table) < 2: print(f"⚠️ Trang {page_idx}: Không phát hiện bảng hợp lệ.") continue header = table[0] rows = table[1:] # Lưu header đầu tiên if saved_header is None: saved_header = header print(f"✅ Trang {page_idx}: Lưu header đầu tiên: {saved_header}") # Nếu trang sau không có header rõ → dùng header cũ if len(header) < len(saved_header) or "REGION" not in header[0]: print(f"↩️ Trang {page_idx}: Không có header rõ ràng, dùng lại header trước.") header = saved_header rows = table else: saved_header = header # cập nhật header hợp lệ if len(rows) == 0: print(f"⚠️ Trang {page_idx}: Không có dữ liệu dưới header.") continue try: df = pd.DataFrame(rows, columns=header) all_dfs.append(df) print(f"✅ Trang {page_idx}: {len(df)} dòng được thêm.") except Exception as e: print(f"❌ Lỗi tạo DataFrame ở trang {page_idx}: {e}") if all_dfs: final_df = pd.concat(all_dfs, ignore_index=True).dropna(how="all").reset_index(drop=True) print(f"✅ Tổng cộng {len(final_df)} dòng được trích xuất từ PDF.") # Xuất ra file tạm (Excel + JSON) base_name = os.path.splitext(filename)[0] tmp_dir = tempfile.gettempdir() # json_path = os.path.join(tmp_dir, f"{base_name}.json") # excel_path = os.path.join(tmp_dir, f"{base_name}.xlsx") # final_df.to_json(json_path, orient="records", force_ascii=False, indent=2) # final_df.to_excel(excel_path, index=False) # print(f"✅ Xuất JSON: {json_path}") # print(f"✅ Xuất Excel: {excel_path}") # Convert bảng thành CSV text để Gemini đọc tiếp table_text = final_df.to_csv(index=False) print(f"✅ Đang Gen text từ file CSV") question = ( f"{PROMPT_FREIGHT_JSON}\n" "Below is the table text extracted from the PDF (CSV format):\n" f"{table_text}\n\n" "Please convert this into valid JSON as per the schema." ) else: print("⚠️ Không có bảng hợp lệ để extract bằng pdfplumber.") except Exception as e: print("❌ pdfplumber extract failed:", e) # STEP 2️⃣: Route model if model_choice == EXTERNAL_MODEL_NAME: return run_process_external( file_bytes=file_bytes, filename=filename, mime=mime, question=question, api_url=external_api_url, temperature=temperature, top_p=top_p ) return run_process_internal_base_v2( file_bytes=file_bytes, filename=filename, mime=mime, question=question, model_choice=model_choice, temperature=temperature, top_p=top_p ) except Exception as e: return f"ERROR: {type(e).__name__}: {str(e)}", None # ================== UI ================== def main(): with gr.Blocks(title="OCR Multi-Agent System") as demo: file = gr.File(label="Upload PDF/Image") question = gr.Textbox(label="Prompt", lines=2) model_choice = gr.Dropdown(choices=[*INTERNAL_MODEL_MAP.keys(), EXTERNAL_MODEL_NAME], value="Gemini 2.5 Flash", label="Model") temperature = gr.Slider(0.0, 2.0, value=0.2, step=0.05) top_p = gr.Slider(0.0, 1.0, value=0.95, step=0.01) external_api_url = gr.Textbox(label="External API URL", visible=False) output_text = gr.Code(label="Output", language="json") run_btn = gr.Button("🚀 Process") run_btn.click( run_process, inputs=[file, question, model_choice, temperature, top_p, external_api_url], outputs=[output_text, gr.State()] ) return demo demo = main() if __name__ == "__main__": demo.launch()