MLBench/logistics_ocr

1import gradio as gr2import json3import os4from pathlib import Path5from typing import List, Dict, Any, Optional6import traceback7 8from PIL import Image9import PyPDF210import pytesseract11from pdf2image import convert_from_path12from huggingface_hub import InferenceClient13 14 15# ==============================================================16# Extraction prompt17# ==============================================================18 19EXTRACTION_PROMPT = """You are an expert shipping-document data extractor.20You will be given OCR/text extracted from shipping documents.21 22Extract and return ONLY valid JSON matching this schema:23 24{25  "poNumber": string | null,26  "shipFrom": string | null,27  "carrierType": string | null,28  "originCarrier": string | null,29  "railCarNumber": string | null,30  "totalQuantity": number | null,31  "totalUnits": string | null,32  "attachments": [string],33  "accountName": string | null,34  "inventories": {35    "items": [36      {37        "quantityShipped": number | null,38        "inventoryUnits": string | null,39        "pcs": number | null,40        "productName": string | null,41        "productCode": string | null,42        "product": {43          "category": number | null,44          "defaultUnits": string | null,45          "unit": string | null,46          "pcs": number | null,47          "mbf": number | null,48          "sf": number | null,49          "pcsHeight": number | null,50          "pcsWidth": number | null,51          "pcsLength": number | null52        },53        "customFields": [string]54      }55    ]56  }57}58 59Return ONLY JSON. No explanation.60"""61 62 63# ==============================================================64# JSON Helpers65# ==============================================================66 67def extract_json(text: str) -> Dict:68    text = text.strip()69 70    if text.startswith("```"):71        text = text.split("\n", 1)[-1]72        text = text.replace("```", "").strip()73 74    start = text.find("{")75    end = text.rfind("}")76 77    if start == -1 or end == -1:78        raise json.JSONDecodeError("No JSON found", text, 0)79 80    return json.loads(text[start:end+1])81 82 83# ==============================================================84# OCR + TEXT EXTRACTION85# ==============================================================86 87def extract_text_from_pdf(pdf_path: str) -> str:88    try:89        with open(pdf_path, "rb") as f:90            reader = PyPDF2.PdfReader(f)91            text = ""92            for page in reader.pages:93                t = page.extract_text()94                if t:95                    text += t + "\n"96            return text97    except Exception as e:98        return f"PDF text error: {e}"99 100 101def ocr_image(img: Image.Image) -> str:102    if img.mode != "RGB":103        img = img.convert("RGB")104    return pytesseract.image_to_string(img)105 106 107def extract_pdf_with_ocr(pdf_path: str) -> str:108    text = extract_text_from_pdf(pdf_path)109 110    if text and len(text) > 50:111        return text112 113    pages = convert_from_path(pdf_path, dpi=250)114    ocr_text = ""115    for p in pages:116        ocr_text += ocr_image(p) + "\n"117 118    return ocr_text119 120 121def process_files(files: List[str]) -> Dict[str, Any]:122    result = {123        "text_content": "",124        "attachments": []125    }126 127    for f in files:128        name = Path(f).name129        ext = Path(f).suffix.lower()130 131        result["attachments"].append(name)132 133        if ext == ".pdf":134            text = extract_pdf_with_ocr(f)135 136        elif ext in [".jpg", ".jpeg", ".png", ".webp"]:137            img = Image.open(f)138            text = ocr_image(img)139 140        elif ext in [".txt", ".csv"]:141            text = open(f, encoding="utf-8", errors="ignore").read()142 143        elif ext in [".doc", ".docx"]:144            import docx145            doc = docx.Document(f)146            text = "\n".join([p.text for p in doc.paragraphs])147 148        else:149            text = ""150 151        result["text_content"] += f"\n\n=== {name} ===\n{text}"152 153    return result154 155 156# ==============================================================157# HF MODEL CALL (Robust: conversational support)158# ==============================================================159 160def extract_with_hf(processed_data: Dict[str, Any]) -> Dict[str, Any]:161    hf_token = os.getenv("HF_TOKEN")162    model = os.getenv("HF_MODEL", "mistralai/Mistral-7B-Instruct-v0.3")163 164    client = InferenceClient(model=model, token=hf_token)165 166    prompt = (167        EXTRACTION_PROMPT168        + "\n\nDOCUMENT TEXT:\n"169        + processed_data["text_content"]170        + "\n\nATTACHMENTS:\n"171        + json.dumps(processed_data["attachments"])172    )173 174    raw = ""175 176    try:177        # FIRST: try conversational (works for Mistral)178        conv = client.conversational(179            {180                "past_user_inputs": [],181                "generated_responses": [],182                "text": prompt,183            }184        )185        raw = conv["generated_text"]186 187    except Exception as e1:188        try:189            # fallback to chat190            resp = client.chat_completion(191                messages=[192                    {"role": "system", "content": "Return strict JSON only."},193                    {"role": "user", "content": prompt}194                ],195                temperature=0.1,196                max_tokens=3000197            )198            raw = resp.choices[0].message.content199 200        except Exception as e2:201            return {202                "success": False,203                "error": f"Model call failed:\n{e1}\n\n{e2}",204                "traceback": traceback.format_exc()205            }206 207    try:208        parsed = extract_json(raw)209        return {210            "success": True,211            "data": parsed,212            "raw": raw213        }214    except Exception as je:215        return {216            "success": False,217            "error": f"JSON parse error: {je}",218            "raw": raw219        }220 221 222# ==============================================================223# MAIN PROCESS224# ==============================================================225 226def process_documents(files):227    if not files:228        return "❌ Upload file", "{}", ""229 230    paths = [f.name if hasattr(f, "name") else f for f in files]231 232    status = "📄 Extracting text...\n"233    processed = process_files(paths)234 235    status += "🤖 Calling HF model...\n"236    result = extract_with_hf(processed)237 238    if result["success"]:239        json_out = json.dumps(result["data"], indent=2)240        return "✅ Success", json_out, json_out241 242    return f"❌ Extraction failed:\n{result['error']}", "{}", result.get("raw", "")243 244 245# ==============================================================246# UI247# ==============================================================248 249with gr.Blocks() as demo:250    gr.Markdown("# 📄 Logistic OCR – Open Source Version")251 252    file_input = gr.File(file_count="multiple")253    btn = gr.Button("🚀 Extract")254    status = gr.Textbox(label="Status")255    json_out = gr.Code(language="json")256    preview = gr.Textbox(label="Preview")257 258    btn.click(259        process_documents,260        inputs=file_input,261        outputs=[status, json_out, preview]262    )263 264demo.launch(server_name="0.0.0.0", server_port=7860)265