Spaces:
Paused
Paused
| 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() | |