138 lines
5.0 KiB
Python
138 lines
5.0 KiB
Python
import os
|
|
import json
|
|
import base64
|
|
import json_repair
|
|
from utils.pdf_util import PDFUtil
|
|
from utils.logger import logger
|
|
from utils.gpt_utils import chat
|
|
|
|
|
|
def get_base64_pdf_image_list(
|
|
pdf_file: str, pdf_page_index_list: list, output_folder: str = None
|
|
) -> dict:
|
|
if pdf_file is None or pdf_file == "" or not os.path.exists(pdf_file):
|
|
logger.error("pdf_file is not provided")
|
|
return None
|
|
pdf_util = PDFUtil(pdf_file)
|
|
if pdf_page_index_list is None or len(pdf_page_index_list) == 0:
|
|
pdf_page_index_list = list(range(pdf_util.get_page_count()))
|
|
if output_folder is not None and len(output_folder) > 0:
|
|
os.makedirs(output_folder, exist_ok=True)
|
|
pdf_image_info = pdf_util.extract_images(
|
|
pdf_page_index_list=pdf_page_index_list, output_folder=output_folder
|
|
)
|
|
return pdf_image_info
|
|
|
|
|
|
def encode_image(image_path: str):
|
|
if image_path is None or len(image_path) == 0 or not os.path.exists(image_path):
|
|
return None
|
|
with open(image_path, "rb") as image_file:
|
|
return base64.b64encode(image_file.read()).decode("utf-8")
|
|
|
|
|
|
def chat_with_image(
|
|
pdf_file: str,
|
|
pdf_page_index_list: list,
|
|
image_instructions_file: str,
|
|
image_folder: str,
|
|
gpt_folder: str,
|
|
):
|
|
if pdf_file is None or pdf_file == "" or not os.path.exists(pdf_file):
|
|
logger.error("pdf_file is not provided")
|
|
return None
|
|
pdf_image_info = get_base64_pdf_image_list(
|
|
pdf_file, pdf_page_index_list, image_folder
|
|
)
|
|
|
|
with open(image_instructions_file, "r", encoding="utf-8") as file:
|
|
image_instructions = file.read()
|
|
os.makedirs(gpt_folder, exist_ok=True)
|
|
pdf_base_name = os.path.basename(pdf_file).replace(".pdf", "")
|
|
response_list = {}
|
|
for page_index, data in pdf_image_info.items():
|
|
logger.info(f"Processing image in page {page_index}")
|
|
image_file = data.get("img_file", None)
|
|
image_base64 = data.get("img_base64", None)
|
|
response, error = chat(prompt=image_instructions, image_base64=image_base64)
|
|
if error:
|
|
logger.error(f"Error in processing image in page {page_index}")
|
|
continue
|
|
try:
|
|
response_json = json.loads(response)
|
|
except:
|
|
response_json = json_repair.loads(response)
|
|
response_json_file = os.path.join(
|
|
gpt_folder, f"{pdf_base_name}_{page_index}.json"
|
|
)
|
|
with open(response_json_file, "w", encoding="utf-8") as file:
|
|
json.dump(response_json, file, indent=4)
|
|
logger.info(f"Response for image in page {page_index}: {response}")
|
|
logger.info("Done")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Table extraction by image
|
|
# pdf_file = r"/data/emea_ar/small_pdf/382366116.pdf"
|
|
# pdf_page_index_list = [29, 35, 71, 77, 83, 89, 97, 103, 112, 121, 130, 140, 195, 250, 305]
|
|
# pdf_file = r"/data/emea_ar/small_pdf/389171486.pdf"
|
|
# pdf_page_index_list = [13]
|
|
# pdf_file = r"/data/emea_ar/small_pdf/402181770.pdf"
|
|
# pdf_page_index_list = [29]
|
|
# image_instructions_file = r"./instructions/table_extraction_image_prompts_v2.txt"
|
|
# image_output_folder = r"/data/emea_ar/small_pdf_image/"
|
|
# gpt_output_folder = r"/data/emea_ar/output/gpt_image_response/table/"
|
|
# chat_with_image(
|
|
# pdf_file,
|
|
# pdf_page_index_list,
|
|
# image_instructions_file,
|
|
# image_output_folder,
|
|
# gpt_output_folder,
|
|
# )
|
|
|
|
# Data extraction by image
|
|
# pdf_file = r"/data/emea_ar/small_pdf/402181770.pdf"
|
|
# pdf_page_index_list = [29]
|
|
pdf_file = r"/data/emea_ar/small_pdf/389171486.pdf"
|
|
pdf_page_index_list = [13]
|
|
image_output_folder = r"/data/emea_ar/small_pdf_image/"
|
|
gpt_output_folder = r"/data/emea_ar/output/gpt_image_response/data/"
|
|
image_instructions_file = r"./instructions/data_extraction_image_prompts.txt"
|
|
chat_with_image(
|
|
pdf_file,
|
|
pdf_page_index_list,
|
|
image_instructions_file,
|
|
image_output_folder,
|
|
gpt_output_folder,
|
|
)
|
|
|
|
|
|
# Text extraction by image
|
|
# pdf_file = r"/data/emea_ar/small_pdf/389171486.pdf"
|
|
# pdf_page_index_list = [13]
|
|
# image_instructions_file = r"./instructions/text_extraction_image_prompts.txt"
|
|
# image_output_folder = r"/data/emea_ar/small_pdf_image/"
|
|
# gpt_output_folder = r"/data/emea_ar/output/gpt_image_response/text/"
|
|
# chat_with_image(
|
|
# pdf_file,
|
|
# pdf_page_index_list,
|
|
# image_instructions_file,
|
|
# image_output_folder,
|
|
# gpt_output_folder,
|
|
# )
|
|
|
|
# pdf_file = r"/data/emea_ar/small_pdf/389171486.pdf"
|
|
# pdf_page_index_list = [13]
|
|
# image_instructions_file = r"./instructions/table_extraction_image_optimize_prompts.txt"
|
|
# image_output_folder = r"/data/emea_ar/small_pdf_image/"
|
|
# gpt_output_folder = r"/data/emea_ar/output/gpt_image_response/optimized_instructions/"
|
|
# chat_with_image(
|
|
# pdf_file,
|
|
# pdf_page_index_list,
|
|
# image_instructions_file,
|
|
# image_output_folder,
|
|
# gpt_output_folder,
|
|
# )
|
|
|
|
|