OCR-pipeline-python / json-V3.py
abdullah-1111's picture
Update json-V3.py
da2b7d6 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 = """
Extract the following fields from the document (in Arabic only) and return them in JSON format:
1. اسم المكلف
2. عنوان المركز الرئيسي
3. المدينة
4. الحي
5. صندوق البريد
6. الرمز البريدي
7. الهاتف
The required JSON format:
{
"taxpayer_name": "...",
"headquarters_address": "...",
"city": "...",
"district": "...",
"po_box": "...",
"postal_code": "...",
"phone": "..."
}
If any field is missing, return it as 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}")