OCR-pipeline-python / json-V2.py
abdullah-1111's picture
Upload 14 files
2db08ff verified
import base64
import json
import re
import requests
import os
# مفتاح Gemini API
API_KEY = "AIzaSyBr2-dUqHDZkk20hlWeEcpWnVVdkq9fqyE"
# المجلد الذي يحتوي صور cr1
cr1_images_folder = r"C:\Users\ASUS\OneDrive - Binder\Desktop\test.orch\classified\CR1"
# مجلد إخراج ملفات JSON
output_json_folder = r"C:\Users\ASUS\OneDrive - Binder\Desktop\test.orch\classified\CR1\cr1_json"
# Ensure output folder exists
os.makedirs(output_json_folder, exist_ok=True)
# Exact same prompt
prompt = """
Please extract the following fields in Arabic and English from the tax document image:
Taxpayer Name
VAT Registration Number
Effective Registration Date
Taxpayer Address
CR/License
Contract/ID No
Tax period
First Filing due date
Return the result in a JSON format with these keys:
en_taxpayer_name, en_vat_registration_number, en_effective_registration_date, en_taxpayer_address,
en_cr_license, en_contract_id_no, en_tax_period, en_first_filing_due_date,
ar_taxpayer_name, ar_vat_registration_number, ar_effective_registration_date, ar_taxpayer_address,
ar_cr_license, ar_contract_id_no, ar_tax_period, ar_first_filing_due_date
If a field is missing, return null.
"""
url = f"https://generativelanguage.googleapis.com/v1/models/gemini-1.5-flash:generateContent?key={API_KEY}"
headers = {"Content-Type": "application/json"}
# Iterate over all images
for image_name in os.listdir(cr1_images_folder):
if not image_name.lower().endswith(('.jpg', '.jpeg', '.png')):
continue
image_path = os.path.join(cr1_images_folder, image_name)
base_name = os.path.splitext(image_name)[0]
output_file = os.path.join(output_json_folder, base_name + ".json")
# Skip if JSON already exists
if os.path.exists(output_file):
print(f"Skipped {image_name} (JSON file already exists)")
continue
# Read image and convert to base64
with open(image_path, "rb") as f:
image_b64 = base64.b64encode(f.read()).decode()
# Send request to Gemini API
data = {
"contents": [
{
"role": "user",
"parts": [
{"text": prompt},
{
"inline_data": {
"mime_type": "image/jpeg",
"data": image_b64
}
}
]
}
]
}
try:
response = requests.post(url, headers=headers, json=data)
response.raise_for_status()
response_text = response.json()['candidates'][0]['content']['parts'][0]['text']
match = re.search(r"```json\s*(\{.*\})\s*```", response_text, re.DOTALL)
if match:
json_text = match.group(1)
result = json.loads(json_text)
with open(output_file, "w", encoding="utf-8") as f:
json.dump(result, f, ensure_ascii=False, indent=2)
print(f"✅ Processed: {image_name}")
else:
print(f"❌ Failed to extract JSON from: {image_name}")
print(response_text)
except Exception as e:
print(f"❌ Error processing image {image_name}: {e}")