File size: 1,473 Bytes
aa93eef 9169bb3 8fe674c aa93eef 8fe674c aa93eef 9169bb3 8fe674c 683701b aa93eef 8fe674c aa93eef 8fe674c 683701b 8fe674c aa93eef 683701b 9169bb3 683701b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | import json
import logging
import os
from google import genai
from google.genai import types
log = logging.getLogger("cardextractor")
# ponytail: key comes from the GEMINI_API_KEY Space secret; fail loud at startup if missing
client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
MODEL = os.environ.get("GEMINI_MODEL", "gemma-4-26b-a4b-it")
PROMPT = """Extract the contact information from this business card.
Return ONLY JSON, no markdown, with exactly these keys:
{"name":"","company":"","designation":"","phone":"","email":"","website":"","address":""}"""
def extract_contact(image_path):
with open(image_path, "rb") as f:
image_bytes = f.read()
# ponytail: let the image content-type ride on the extension; default jpeg
ext = os.path.splitext(image_path)[1].lower().lstrip(".") or "jpeg"
mime = "image/png" if ext == "png" else f"image/{ext}"
# ponytail: Gemma models don't support response_mime_type=json, so parse
# the JSON object out of the text ourselves (handles ```json fences too)
resp = client.models.generate_content(
model=MODEL,
contents=[
types.Part.from_bytes(data=image_bytes, mime_type=mime),
PROMPT,
],
)
text = resp.text
log.info("LLM response: %s", text)
start = text.find("{")
if start == -1:
raise ValueError(f"no JSON in model output: {text!r}")
data, _ = json.JSONDecoder().raw_decode(text[start:])
return data
|