CoolFace
Apppublic

aibridze/document_intelligence

sourceHugging Facemitupdated 6mo agoView on Hugging Face
0likes
vision_model.py125 linesDownload Raw Back to core
1from google import genai2from google.genai import types3from typing import List4import base645import asyncio6from app.config import get_settings7 8settings = get_settings()9 10 11class GeminiVision:12    13    def __init__(self):14        if not settings.GOOGLE_API_KEY:15            raise ValueError("GOOGLE_API_KEY not configured. Please set it in .env file.")16        17        self.client = genai.Client(api_key=settings.GOOGLE_API_KEY)18        self.model_name = settings.VISION_MODEL19    20    async def analyze_document(self, images: List[str], prompt: str) -> str:21        22        contents = []23        24        # Add images25        for img_b64 in images:26            img_bytes = base64.b64decode(img_b64)27            contents.append(types.Part.from_bytes(data=img_bytes, mime_type="image/png"))28        29        # Add prompt30        contents.append(prompt)31        32        # Generate response (run sync call in thread to avoid blocking event loop)33        response = await asyncio.to_thread(34            self.client.models.generate_content,35            model=self.model_name,36            contents=contents,37            config=types.GenerateContentConfig(38                temperature=0,39                max_output_tokens=819240            )41        )42        43        return response.text44    45    async def extract_handwriting(self, images: List[str]) -> str:46        prompt = """47        Analyze these document pages carefully and:48        49        1. Identify ALL handwritten notes, signatures, or modifications50        2. Transcribe handwritten text EXACTLY as written51        3. Note which printed text the handwriting modifies or annotates52        4. Include location descriptions (e.g., "margin of page 1", "bottom of clause 5")53        54        Return a structured list in this format:55        56        HANDWRITTEN CONTENT FOUND:57        58        1. Location: [where in the document]59           Text: [transcribed handwritten text]60           Modifies: [what printed text this relates to, if any]61        62        2. [next item...]63        64        If no handwriting is detected, return: "No handwritten content detected."65        """66        67        return await self.analyze_document(images, prompt)68    69    async def extract_all_text(self, images: List[str]) -> str:70        71        prompt = """You are an expert document OCR system. Extract ALL text from these document pages with maximum accuracy.72 73CRITICAL INSTRUCTIONS:741. Extract EVERY piece of text — printed, typed, handwritten, stamped, or watermarked.752. Preserve the document structure as faithfully as possible:76   - Keep numbered clauses and sub-clauses in order (e.g., 1, 1.1, 1.1.1)77   - Preserve table structures — use a clear text format for tables78   - Keep headers, footers, and section titles79   - Preserve paragraph breaks803. Mark each page clearly with "--- Page X ---" headers.814. For tables, use pipe-delimited format:82   | Column1 | Column2 | Column3 |83   | data    | data    | data    |845. For handwritten text, include it inline with [HANDWRITTEN: text] markers.856. For signatures, note them as [SIGNATURE: name if legible]867. For stamps/seals, note them as [STAMP: text if legible]878. Include ALL annexures, schedules, and appendices.889. Do NOT summarize or paraphrase — extract the EXACT text.89"""90        91        return await self.analyze_document(images, prompt)92    93    async def analyze_contract_structure(self, images: List[str]) -> str:94        95        prompt = """96        Analyze this contract document and identify:97        98        1. CONTRACT TYPE (NDA, Service Agreement, Work Order, Supply Agreement, Consultancy Agreement, etc.)99        100        2. PARTIES INVOLVED:101           - Full name of each party102           - Role/designation in contract (e.g., "Licensor", "Service Provider", "Supplier", "Buyer")103           - Address if visible104        105        3. KEY DATES:106           - Execution/signing date107           - Effective date108           - Expiry/termination date109        110        4. DEFINED TERMS:111           - List all defined terms (e.g., "Borrower", "Licensee") and their actual values112        113        5. SECTION STRUCTURE:114           - List main sections/clauses visible115        116        Return in structured format.117        """118        119        return await self.analyze_document(images, prompt)120 121 122def get_vision_model() -> GeminiVision:123    """Get configured Gemini vision model instance."""124    return GeminiVision()125