ReceiptSplitAI / process_images_gemini_one_hread.py
valentynliubchenko
merging
eba303d
import os
import base64
import json
import time
from utils import encode_image_to_jpeg_base64
# Import your VertexAIService class
from vertex_ai_service import VertexAIService # Replace with the module name where VertexAIService is defined
def process_image_file(file_path, client):
"""
Processes an image file and returns the path to the JSON file with results.
:param file_path: Path to the image file
:param client: Instance of VertexAIService
:return: Path to the saved JSON file
"""
input_image64 = encode_image_to_jpeg_base64(file_path)
# Processing settings
prompt = "Read the text"
system = "You are receipt recognizer"
# Call image processing
result_img = client.process_image(input_image64, "gemini-1.5-flash", prompt, system, 0.0)
# Create path for JSON file
json_file_path = os.path.splitext(file_path)[0] + ".json"
# Write the result to a JSON file
with open(json_file_path, 'w', encoding='utf-8') as json_file:
json_file.write(result_img)
return json_file_path
def process_folder_images_sequentially(folder_path, json_key_path):
"""
Processes all images in the specified folder sequentially.
:param folder_path: Path to the folder containing images
:param json_key_path: Path to the JSON key for Vertex AI authentication
"""
# Initialize Vertex AI client
client = VertexAIService(json_key_path=json_key_path)
# Get a list of all images in the folder
image_files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith(('.jpg', '.jpeg', '.png', 'webp'))]
# Process each image file sequentially
for file_path in image_files:
try:
json_file_path = process_image_file(file_path, client)
print(f"Processed: {file_path}, result saved to: {json_file_path}")
time.sleep(60)
except Exception as exc:
print(f"Failed to process file {file_path}. Error: {exc}")
# Call the function to process the folder
if __name__ == '__main__':
folder_path = './examples' # Set the path to your folder
json_key_path = 'secrets/GOOGLE_VERTEX_AI_KEY.json' # Set the path to your JSON key
process_folder_images_sequentially(folder_path, json_key_path)