Create app.py

app.py

@@ -0,0 +1,408 @@
+from pdf2image import convert_from_path
+from PIL import Image
+from pathlib import Path
+from torchvision.transforms.functional import InterpolationMode
+from transformers import AutoModel, AutoTokenizer
+from pydantic import BaseModel, Field, ValidationError, root_validator, validator
+from typing import List, Optional, Literal
+
+import gradio as gr
+import tempfile
+import os
+import json
+import numpy as np
+import re
+import torch
+import torchvision.transforms as T
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+def build_transform(input_size):
+    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
+    transform = T.Compose([
+        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
+        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
+        T.ToTensor(),
+        T.Normalize(mean=MEAN, std=STD)
+    ])
+    return transform
+
+def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
+    best_ratio_diff = float('inf')
+    best_ratio = (1, 1)
+    area = width * height
+    for ratio in target_ratios:
+        target_aspect_ratio = ratio[0] / ratio[1]
+        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+        if ratio_diff < best_ratio_diff:
+            best_ratio_diff = ratio_diff
+            best_ratio = ratio
+        elif ratio_diff == best_ratio_diff:
+            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+                best_ratio = ratio
+    return best_ratio
+
+def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
+    orig_width, orig_height = image.size
+    aspect_ratio = orig_width / orig_height
+
+    # calculate the existing image aspect ratio
+    target_ratios = set(
+        (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
+        i * j <= max_num and i * j >= min_num)
+    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+    # find the closest aspect ratio to the target
+    target_aspect_ratio = find_closest_aspect_ratio(
+        aspect_ratio, target_ratios, orig_width, orig_height, image_size)
+
+    # calculate the target width and height
+    target_width = image_size * target_aspect_ratio[0]
+    target_height = image_size * target_aspect_ratio[1]
+    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+    # resize the image
+    resized_img = image.resize((target_width, target_height))
+    processed_images = []
+    for i in range(blocks):
+        box = (
+            (i % (target_width // image_size)) * image_size,
+            (i // (target_width // image_size)) * image_size,
+            ((i % (target_width // image_size)) + 1) * image_size,
+            ((i // (target_width // image_size)) + 1) * image_size
+        )
+        # split the image
+        split_img = resized_img.crop(box)
+        processed_images.append(split_img)
+    assert len(processed_images) == blocks
+    if use_thumbnail and len(processed_images) != 1:
+        thumbnail_img = image.resize((image_size, image_size))
+        processed_images.append(thumbnail_img)
+    return processed_images
+
+def load_image(image_file, input_size=448, max_num=12):
+    image = Image.open(image_file).convert('RGB')
+    transform = build_transform(input_size=input_size)
+    images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
+    pixel_values = [transform(image) for image in images]
+    pixel_values = torch.stack(pixel_values)
+    return pixel_values
+
+device = torch.device('cpu')
+base_path = Path().absolute()
+path = base_path / "resources/model/InternVL2-1B"
+model = AutoModel.from_pretrained(
+    path,
+    low_cpu_mem_usage=True,
+    trust_remote_code=True).to(device).eval()
+tokenizer = AutoTokenizer.from_pretrained(str(path), trust_remote_code=True, use_fast=False)
+
+def convert_pdf_to_images(pdf_path):
+    images = convert_from_path(pdf_path)
+    temp_file_paths = []
+    for i, image in enumerate(images):
+        with tempfile.NamedTemporaryFile(delete=False, suffix=f'_{i+1}.png', mode='wb') as temp_file:
+            temp_path = temp_file.name
+            image.save(temp_path, format='PNG')
+            temp_file_paths.append(temp_path)
+    return temp_file_paths
+
+
+def images_to_pixel_values(image_paths, device):
+    image_paths.sort(key=lambda x: int(os.path.splitext(os.path.basename(x))[0].split('_')[-1]))
+    pixel_values_list = [load_image(img_path, max_num=12).to(device) for img_path in image_paths]
+    return torch.cat(pixel_values_list, dim=0)
+
+
+def process_file(file, file_type, file_document):
+    if file_type == "PDF":
+        try:
+            images = convert_pdf_to_images(file.name)
+            pixel_values = images_to_pixel_values(images, device)
+
+            if file_document == "Receipt":
+                question = '''
+                <image> You are a document processing model. This is a purchase receipt. If multiple images are provided, treat them as a single document and combine their content for processing. Extract and label the following entities from the given image, providing the output in JSON format. The values must only include text found in the document image. Use null or [] for missing entity types.
+                { "type": "object", "properties": { "document_type": { "type": "string", "enum": ["purchase_receipt"] }, "data": { "type": "object", "properties": { "receipt_number": { "type": "string" }, "vendor_name": { "type": "string" }, "customer_name": { "type": "string" }, "items": { "type": "array", "items": { "type": "object", "properties": { "description": { "type": "string" }, "quantity": { "type": "integer", "nullable": true }, "unit_price": { "type": "string" }, "total_price": { "type": "string" } }, "required": ["description", "quantity", "unit_price", "total_price"] } }, "total_amount": { "type": "string" } }, "required": ["receipt_number", "vendor_name", "customer_name", "items", "total_amount"] } }, "required": ["document_type", "data"] }
+                '''
+                generation_config = dict(num_beams=1,
+                                         max_new_tokens=4096,
+                                         do_sample=True,
+                                         temperature=0.2,
+                                         repetition_penalty=1.1,
+                                         top_p=0.7)
+                response = model.chat(tokenizer, pixel_values, question, generation_config)
+
+            elif file_document == "Invoice":
+                question = '''
+                <image> You are a document processing model. This is an invoice. If multiple images are provided, treat them as a single document and combine their content for processing. Extract and label the following entities from the given image, providing the output in JSON format. The values must only include text found in the document image. Use null or [] for missing entity types. 
+                { "type": "object", "properties": { "document_type": { "type": "string", "enum": ["invoice"] }, "data": { "type": "object", "properties": { "invoice_number": { "type": "string" }, "vendor_name": { "type": "string" }, "customer_name": { "type": "string" }, "items": { "type": "array", "items": { "type": "object", "properties": { "description": { "type": "string" }, "quantity": { "type": "integer", "nullable": true }, "unit_price": { "type": "string" }, "total_price": { "type": "string" } }, "required": ["description", "quantity", "unit_price", "total_price"] } }, "subtotal": { "type": "string" }, "tax": { "type": "string" }, "total_amount": { "type": "string" } }, "required": ["invoice_number", "vendor_name", "customer_name", "items", "subtotal", "tax", "total_amount"] } }, "required": ["document_type", "data"] }
+                '''
+                generation_config = dict(max_new_tokens=4096)
+                response = model.chat(tokenizer, pixel_values, question, generation_config)
+
+            elif file_document == "Faktur Pajak":
+                question = '''
+                <image> You are a document processing model. This is an invoice. If multiple images are provided, treat them as a single document and combine their content for processing. Extract and label the following entities from the given image, providing the output in JSON format. The values must only include text found in the document image. Use null or [] for missing entity types. 
+                { "type": "object", "properties": { "document_type": { "type": "string", "enum": ["invoice"] }, "data": { "type": "object", "properties": { "invoice_number": { "type": "string" }, "vendor_name": { "type": "string" }, "customer_name": { "type": "string" }, "items": { "type": "array", "items": { "type": "object", "properties": { "description": { "type": "string" }, "quantity": { "type": "integer", "nullable": true }, "unit_price": { "type": "string" }, "total_price": { "type": "string" } }, "required": ["description", "quantity", "unit_price", "total_price"] } }, "subtotal": { "type": "string" }, "tax": { "type": "string" }, "total_amount": { "type": "string" } }, "required": ["invoice_number", "vendor_name", "customer_name", "items", "subtotal", "tax", "total_amount"] } }, "required": ["document_type", "data"] }
+                '''
+                generation_config = dict(max_new_tokens=4096)
+                response = model.chat(tokenizer, pixel_values, question, generation_config)
+
+            elif file_document == "E-Statement":
+               question = '''
+               <image> You are a document processing model. This is an e-statement. If multiple images are provided, treat them as a single document and combine their content for processing. Extract and label the following entities from the given image, providing the output in JSON format. The values must only include text found in the document image. Use null or [] for missing entity types.
+               { "type": "object", "properties": { "document_type": { "type": "string", "enum": ["e-statement"] }, "data": { "type": "object", "properties": { "account_number": { "type": "string" }, "bank_name": { "type": "string" }, "customer_name": { "type": "string" }, "statement_period": { "type": "string" }, "currency": { "type": "string" }, "country": { "type": "string" }, "transactions": { "type": "array", "items": { "type": "object", "properties": { "transaction_type": { "type": "string" }, "amount": { "type": "number", "nullable": true }, "date": { "type": "string" }, "reference": { "type": "string" } }, "required": ["transaction_type", "amount", "date", "reference"] } } }, "required": ["account_number", "bank_name", "customer_name", "statement_period", "currency", "country", "transactions"] } }, "required": ["document_type", "data"] }
+               '''
+               generation_config = dict(num_beams=1,
+                                        max_new_tokens=4096,
+                                        do_sample=True,
+                                        repetition_penalty=1.1,
+                                        temperature=0.2,
+                                        top_p=0.7)
+               response = model.chat(tokenizer, pixel_values, question, generation_config)
+
+            elif file_document == "Hukum":
+               question = "You are a document processing model. The provided documents may include laws, regulations, legal journals, or other related legal papers. Extract and label the relevant entities from the given documents. If multiple images or files are provided, treat them as a single document and combine their content for processing. Provide the output in JSON format."
+               generation_config = dict(max_new_tokens=4096,
+                                        repetition_penalty=1.2,
+                                        do_sample=True,
+                                        temperature=0.1,
+                                        top_p=1.0)
+               response = model.chat(tokenizer, pixel_values, question, generation_config)
+
+        except Exception as e:
+            return None, f"Terjadi kesalahan saat memproses file PDF: {str(e)}"
+
+    elif file_type == "Image":
+        try:
+            image = Image.open(file.name)
+            
+            with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file:
+                temp_path = temp_file.name
+                image.save(temp_path, format='PNG')
+
+            pixel_values = load_image(temp_path, max_num=12).to(device)
+
+            if file_document == "Receipt":
+                question = '''
+                <image> You are a document processing model. This is a purchase receipt. Extract and label the following entities from the given image, providing the output in JSON format. The values must only include text found in the document image. Use null or [] for missing entity types. Extra The output JSON format must follow these specifications:
+                { "type": "object", "properties": { "document_type": { "type": "string", "enum": ["purchase_receipt"] }, "data": { "type": "object", "properties": { "receipt_number": { "type": "string" }, "vendor_name": { "type": "string" }, "customer_name": { "type": "string" }, "items": { "type": "array", "items": { "type": "object", "properties": { "description": { "type": "string" }, "quantity": { "type": "integer", "nullable": true }, "unit_price": { "type": "string" }, "total_price": { "type": "string" } }, "required": ["description", "quantity", "unit_price", "total_price"] } }, "total_amount": { "type": "string" } }, "required": ["receipt_number", "vendor_name", "customer_name", "items", "total_amount"] } }, "required": ["document_type", "data"] }
+                '''
+                generation_config = dict(num_beams=1,
+                                         max_new_tokens=4096,
+                                         do_sample=True,
+                                         temperature=0.2,
+                                         repetition_penalty=1.1,
+                                         top_p=0.7)
+                response = model.chat(tokenizer, pixel_values, question, generation_config)
+
+            elif file_document == "Invoice":
+                question = '''
+                <image> You are a document processing model. This is an invoice. Extract and label the following entities from the given image, providing the output in JSON format. The values must only include text found in the document image. Use null or [] for missing entity types. The output JSON format must follow these specifications:
+                { "type": "object", "properties": { "document_type": { "type": "string", "enum": ["invoice"] }, "data": { "type": "object", "properties": { "invoice_number": { "type": "string" }, "vendor_name": { "type": "string" }, "customer_name": { "type": "string" }, "items": { "type": "array", "items": { "type": "object", "properties": { "description": { "type": "string" }, "quantity": { "type": "integer", "nullable": true }, "unit_price": { "type": "string" }, "total_price": { "type": "string" } }, "required": ["description", "quantity", "unit_price", "total_price"] } }, "subtotal": { "type": "string" }, "tax": { "type": "string" }, "total_amount": { "type": "string" } }, "required": ["invoice_number", "vendor_name", "customer_name", "items", "subtotal", "tax", "total_amount"] } }, "required": ["document_type", "data"] }
+                '''
+                generation_config = dict(max_new_tokens=4096)
+                response = model.chat(tokenizer, pixel_values, question, generation_config)
+
+            elif file_document == "Faktur Pajak":
+                question = '''
+                <image> You are a document processing model. This is an invoice. Extract and label the following entities from the given image, providing the output in JSON format. The values must only include text found in the document image. Use null or [] for missing entity types. 
+                { "type": "object", "properties": { "document_type": { "type": "string", "enum": ["invoice"] }, "data": { "type": "object", "properties": { "invoice_number": { "type": "string" }, "vendor_name": { "type": "string" }, "customer_name": { "type": "string" }, "items": { "type": "array", "items": { "type": "object", "properties": { "description": { "type": "string" }, "quantity": { "type": "integer", "nullable": true }, "unit_price": { "type": "string" }, "total_price": { "type": "string" } }, "required": ["description", "quantity", "unit_price", "total_price"] } }, "subtotal": { "type": "string" }, "tax": { "type": "string" }, "total_amount": { "type": "string" } }, "required": ["invoice_number", "vendor_name", "customer_name", "items", "subtotal", "tax", "total_amount"] } }, "required": ["document_type", "data"] }
+                '''
+                generation_config = dict(max_new_tokens=4096)
+                response = model.chat(tokenizer, pixel_values, question, generation_config)
+
+            elif file_document == "E-Statement":
+               question = '''
+               <image> You are a document processing model. This is an e-statement. Extract and label the following entities from the given image, providing the output in JSON format. The values must only include text found in the document image. Use null or [] for missing entity types.
+               { "type": "object", "properties": { "document_type": { "type": "string", "enum": ["e-statement"] }, "data": { "type": "object", "properties": { "account_number": { "type": "string" }, "bank_name": { "type": "string" }, "customer_name": { "type": "string" }, "statement_period": { "type": "string" }, "currency": { "type": "string" }, "country": { "type": "string" }, "transactions": { "type": "array", "items": { "type": "object", "properties": { "transaction_type": { "type": "string" }, "amount": { "type": "number", "nullable": true }, "date": { "type": "string" }, "reference": { "type": "string" } }, "required": ["transaction_type", "amount", "date", "reference"] } } }, "required": ["account_number", "bank_name", "customer_name", "statement_period", "currency", "country", "transactions"] } }, "required": ["document_type", "data"] }
+               '''
+               generation_config = dict(num_beams=1,
+                                        max_new_tokens=4096,
+                                        do_sample=True,
+                                        repetition_penalty=1.1,
+                                        temperature=0.2,
+                                        top_p=0.7)
+               response = model.chat(tokenizer, pixel_values, question, generation_config)
+
+            elif file_document == "Hukum":
+               question = "You are a document processing model. The provided documents may include laws, regulations, legal journals, or other related legal papers. Extract and label the relevant entities from the given documents. If multiple images or files are provided, treat them as a single document and combine their content for processing. Provide the output in JSON format."
+               generation_config = dict(max_new_tokens=4096,
+                                        repetition_penalty=1.2,
+                                        do_sample=True,
+                                        temperature=0.1,
+                                        top_p=1.0)
+               response = model.chat(tokenizer, pixel_values, question, generation_config)
+            
+
+        except Exception as e:
+            return None, f"Terjadi kesalahan saat memproses file gambar: {str(e)}"
+
+    def preprocess_json_string_for_numbers(json_string: str) -> str:
+        def replace_misplaced_separators(match):
+            num_str = match.group(0)
+            num_str = re.sub(r'[^\d,.-]', '', num_str)
+            
+            if ',' in num_str and '.' not in num_str:
+                corrected_num_str = num_str.replace(',', '')
+            elif ',' in num_str and '.' in num_str:
+                if num_str.index(',') > num_str.index('.'):
+                    corrected_num_str = num_str.replace('.', '').replace(',', '.')
+                else:
+                    corrected_num_str = num_str.replace(',', '')
+            else:
+                corrected_num_str = num_str.replace('.', '')
+    
+            corrected_num_str = str(int(corrected_num_str)) 
+            return corrected_num_str
+    
+        corrected_string = re.sub(r'\b\d{1,3}([.,]\d{3})*(,\d{2})?\b', replace_misplaced_separators, json_string)
+        return corrected_string
+
+    def extract_and_parse_json(string_with_json: str):
+        json_pattern = re.search(r'{.*}', string_with_json, re.DOTALL)
+        
+        if json_pattern:
+            json_string = json_pattern.group()
+            
+            try:
+                corrected_json_string = preprocess_json_string_for_numbers(json_string)
+                json_data = json.loads(corrected_json_string)
+                print("Parsed JSON data:", json_data)
+                return json_data
+            except json.JSONDecodeError as e:
+                print("Error decoding JSON:", e)
+                return corrected_json_string
+        else:
+            print("No JSON found in the string")
+            return None
+
+    string_with_json = response
+    string_with_json = string_with_json.replace("'", '"')
+    parsed_data = extract_and_parse_json(string_with_json)
+
+    if file_document == "Receipt":
+        class Item(BaseModel):
+            description: Optional[str] = None
+            quantity: Optional[int] = 1
+            unit_price: Optional[float] = None
+            total_price: Optional[float] = None
+        
+            @validator('unit_price', 'total_price', pre=True, always=True)
+            def parse_price(cls, v):
+                if isinstance(v, str):
+                    v = re.sub(r'[^\d,.-]', '', v).replace(',', '')
+                    return float(v)
+                return v
+
+        class PurchaseReceiptData(BaseModel):
+            receipt_number: Optional[str] = None
+            vendor_name: Optional[str] = None
+            customer_name: Optional[str] = None
+            items: List[Item]
+            total_amount: Optional[float] = None
+        
+            @validator('total_amount', pre=True, always=True)
+            def remove_currency_format(cls, v):
+                if isinstance(v, str):
+                    v = re.sub(r'[^\d,.-]', '', v).replace(',', '')
+                    return float(v)
+                return v
+        
+        class PurchaseReceipt(BaseModel):
+            document_type: Literal["purchase_receipt"]
+            data: PurchaseReceiptData
+        
+            @root_validator(pre=True)
+            def ensure_format(cls, values):
+                if 'document_type' not in values:
+                    values['document_type'] = 'purchase_receipt'
+                return values
+        
+        predicted_json = parsed_data
+        
+        try:
+            receipt = PurchaseReceipt(document_type="purchase_receipt", data=predicted_json['data'])
+            corrected_json = receipt.dict()
+            print(corrected_json)
+        except ValidationError as e:
+            print(f"Validation error: {e}")
+
+    elif file_document == "Invoice":
+        class Item(BaseModel):
+            description: Optional[str] = None
+            quantity: Optional[int] = 1
+            unit_price: Optional[float] = None
+            total_price: Optional[float] = None
+
+            @validator('unit_price', 'total_price', pre=True, always=True)
+            def parse_price(cls, v):
+                if isinstance(v, str):
+                    v = re.sub(r'[^\d,.-]', '', v).replace(',', '')
+                    return float(v)
+                return v
+        
+        class InvoiceData(BaseModel):
+            invoice_number: Optional[str] = None
+            vendor_name: Optional[str] = None
+            customer_name: Optional[str] = None
+            items: List[Item]
+            subtotal: float
+            tax: Optional[float] = None
+            total_amount: float
+
+            @validator('total_amount', 'subtotal', 'tax', pre=True, always=True)
+            def remove_currency_format(cls, v):
+                if isinstance(v, str):
+                    v = re.sub(r'[^\d,.-]', '', v).replace(',', '')
+                    return float(v)
+                return v
+        
+        class Invoice(BaseModel):
+            document_type: Literal["invoice"]
+            data: InvoiceData
+        
+            @root_validator(pre=True)
+            def ensure_format(cls, values):
+                if 'document_type' not in values:
+                    values['document_type'] = 'invoice'
+                return values
+        
+        predicted_json = parsed_data
+        
+        try:
+            invoice = Invoice(document_type="invoice", data=predicted_json['data'])
+            corrected_json = invoice.dict()
+            print(corrected_json)
+        except ValidationError as e:
+            print(f"Validation error: {e}")
+
+    elif file_document == "Faktur Pajak":
+        predicted_json = parsed_data
+        corrected_json = predicted_json
+
+    elif file_document == "E-Statement":
+        predicted_json = parsed_data
+        corrected_json = predicted_json
+
+    elif file_document == "Hukum":
+            predicted_json = parsed_data
+            corrected_json = predicted_json
+
+    return corrected_json
+
+file_type_radio = gr.Radio(choices=["Image", "PDF"], label="Pilih tipe file")
+file_document_radio = gr.Radio(choices=["Receipt", "Invoice", "Faktur Pajak", "E-Statement", "Hukum"], label="Pilih jenis dokumen file")
+file_input = gr.File(label="Unggah file", file_types=["image", "pdf"])
+
+interface = gr.Interface(
+    fn=process_file,
+    inputs=[file_input, file_type_radio, file_document_radio],
+    outputs="json",
+    title="POC Finance Document Processing Using AI",
+    description="Pilih tipe file dan unggah file gambar atau PDF untuk diproses."
+)
+
+interface.launch()