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 = ''' 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 = ''' 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 = ''' 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 = ''' 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 = ''' 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 = ''' 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 = ''' 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 = ''' 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()