app.py

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')
path = "./model"
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()