Spaces:
Runtime error
Runtime error
File size: 2,578 Bytes
fa78a29 | 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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import io
import json
import os
from PIL import Image
from google import genai
from google.genai import types
from dotenv import load_dotenv
load_dotenv()
# Models tried in order; first successful response wins.
MODELS = [
"gemini-3-flash-preview",
"gemini-2.5-flash",
"gemini-2.0-flash",
]
_client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
_JSON_CONFIG = types.GenerateContentConfig(
response_mime_type="application/json",
thinking_config=types.ThinkingConfig(thinking_budget=0),
)
def smart_resize(image_bytes: bytes, threshold: int = 1400) -> bytes:
"""Downscale images larger than `threshold` px on either side to 1024×1024."""
img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
if img.width > threshold or img.height > threshold:
img.thumbnail((1024, 1024))
buf = io.BytesIO()
img.save(buf, format="JPEG", quality=85)
return buf.getvalue()
return image_bytes
class GeminiExtractor:
"""
Generic Gemini JSON extractor.
Sends an image + prompt to Gemini and returns parsed JSON.
Falls back across `models` on failure, raising RuntimeError if all fail.
Args:
prompt: Instruction prompt sent alongside the image.
models: Model list to try in order (defaults to module-level MODELS).
resize: Pre-shrink large images before sending (default True).
Set False when the returned coordinates must align with the
original image dimensions (e.g. bounding-box detection).
"""
def __init__(
self,
prompt: str,
models: list[str] | None = None,
resize: bool = True,
) -> None:
self._prompt = prompt
self._models = models or MODELS
self._resize = resize
def extract(self, image_bytes: bytes) -> dict:
"""Return ``{"engine_used": str, "data": dict}``."""
if self._resize:
image_bytes = smart_resize(image_bytes)
image_part = types.Part.from_bytes(data=image_bytes, mime_type="image/jpeg")
last_error = None
for model in self._models:
try:
response = _client.models.generate_content(
model=model,
contents=[image_part, self._prompt],
config=_JSON_CONFIG,
)
return {"engine_used": model, "data": json.loads(response.text)}
except Exception as e:
last_error = str(e)
raise RuntimeError(f"All models failed. Last error: {last_error}")
|