Spaces:
Sleeping
Sleeping
Download gemini_client.py from rairo/chilltrack-api: direct link, hf CLI and curl.
- Browser
- Download file 10.9 kB
-
https://huggingface.co/spaces/rairo/chilltrack-api/resolve/main/gemini_client.py
- Command line
-
hf download hf://spaces/rairo/chilltrack-api/gemini_client.py
-
curl -L -o gemini_client.py https://huggingface.co/spaces/rairo/chilltrack-api/resolve/main/gemini_client.py
10.9 kB
| """ | |
| ChillTrack API — Gemini vision extraction for chiller/freezer gauge readings. | |
| Follows the google-genai SDK usage pattern from NeoFix-API/main.py and | |
| iristrack-api/main.py: `from google import genai`, `genai.Client(api_key=...)`, | |
| `types.Part.from_bytes(...)` for image bytes, `response_mime_type='application/json'` | |
| plus `response_schema=...` for structured output, and parsing the response by | |
| joining `resp.candidates[0].content.parts`. | |
| The Gemini API key env var is literally named `Gemini` (not GEMINI_API_KEY), | |
| matching every sibling repo. | |
| """ | |
| import os | |
| import json | |
| from google import genai | |
| from google.genai import types | |
| MODEL_NAME = "gemini-3.1-flash-lite" | |
| PROMPT_VERSION = "v1" | |
| # response_schema is an OpenAPI 3.0 subset — no additionalProperties, no $ref. | |
| READING_SCHEMA = { | |
| "type": "object", | |
| "properties": { | |
| "value_c": { | |
| "type": "number", | |
| "description": "The air/chamber temperature reading in Celsius, read from the gauge or LCD.", | |
| }, | |
| "confidence": { | |
| "type": "number", | |
| "description": "Confidence in the reading, 0.0 to 1.0.", | |
| }, | |
| "detected_display_type": { | |
| "type": "string", | |
| "enum": ["analog_dial", "digital_lcd", "unclear"], | |
| }, | |
| "raw_text_seen": { | |
| "type": "string", | |
| "description": "Any raw text/digits actually visible on the display, verbatim.", | |
| }, | |
| "setpoint_text_seen": { | |
| "type": "string", | |
| "nullable": True, | |
| "description": "The SETPOINT/Code field text if a digital controller shows one, else null.", | |
| }, | |
| "image_quality": { | |
| "type": "string", | |
| "enum": ["ok", "blurry", "glare", "dusty", "too_dark", "obstructed"], | |
| }, | |
| "needle_angle_estimate_deg": { | |
| "type": "number", | |
| "nullable": True, | |
| "description": "Estimated needle angle in degrees for an analog dial, else null.", | |
| }, | |
| "notes": { | |
| "type": "string", | |
| "description": "Any other observations useful for a human reviewing this reading.", | |
| }, | |
| }, | |
| "required": [ | |
| "value_c", | |
| "confidence", | |
| "detected_display_type", | |
| "raw_text_seen", | |
| "image_quality", | |
| "notes", | |
| ], | |
| } | |
| _client = None | |
| def get_client(): | |
| """Returns a cached google-genai Client, built from the `Gemini` env var.""" | |
| global _client | |
| if _client is not None: | |
| return _client | |
| api_key = os.environ.get("Gemini") | |
| if not api_key: | |
| print("FATAL: The 'Gemini' environment variable for the API key is not set.") | |
| exit(1) | |
| try: | |
| _client = genai.Client(api_key=api_key) | |
| print("Google GenAI Client initialized successfully.") | |
| except Exception as e: | |
| print(f"FATAL: Error initializing GenAI Client: {e}") | |
| exit(1) | |
| return _client | |
| def build_extraction_prompt(unit: dict) -> str: | |
| """Builds the vision-extraction prompt for one unit's photo. | |
| `unit` is a ctk_units doc: {name, setpoint_c, tolerance_c, display_type, ...}. | |
| """ | |
| setpoint_c = unit.get("setpoint_c") | |
| tolerance_c = unit.get("tolerance_c") | |
| name = unit.get("name", "this unit") | |
| return f"""You are reading a temperature gauge photographed by a meat-plant technician | |
| for cold-chain compliance logging. The unit being photographed is: "{name}". | |
| This plant has two kinds of displays. Identify which one is in the photo and read it | |
| according to its type: | |
| 1. ANALOG DIAL GAUGE: a round dial with a needle pointing at a number on a | |
| -40 to +40 degrees Celsius scale. Read the exact number the needle points at | |
| (interpolate between marked graduations if needed). Report your estimate of the | |
| needle's angle in degrees if you can. | |
| 2. DIGITAL LCD REEFER CONTROLLER: an electronic panel with TWO separate fields — | |
| a "SETPOINT" or "Code" field on the LEFT, and an "AIR TEMPERATURE" or "Data" field | |
| on the RIGHT. You must extract ONLY the right-hand air-temperature field as the | |
| reading. Do NOT report the setpoint/code field as the reading — but do report what | |
| it shows in `setpoint_text_seen` so a human can cross-check it. | |
| This unit's expected setpoint is around {setpoint_c} degrees Celsius, +/- {tolerance_c} | |
| degrees Celsius tolerance. Use this ONLY as a sanity check on which field/needle you are | |
| reading — report what you actually see even if it differs from this expectation. Do not | |
| force the reading to match the expected setpoint. | |
| If the image is dusty, glared, poorly lit, or partially obscured, still give your best | |
| numeric estimate. Do not refuse to answer. Instead, lower `confidence` accordingly and | |
| set `image_quality` to describe the problem (blurry / glare / dusty / too_dark / obstructed). | |
| Only use `image_quality: "ok"` when the display is clearly and fully readable. | |
| Return your reading as structured JSON matching the provided schema.""" | |
| # response_schema for a whole handwritten paper form: one photo can cover dozens | |
| # of units across all 12 slot columns. Handwriting OCR at that density is | |
| # inherently less reliable than a single close-up gauge photo, which is exactly | |
| # why every extracted row goes through an admin review/correct screen before | |
| # anything is written to ctk_readings — this schema is a first-pass draft, not | |
| # a final record. | |
| FORM_SCHEMA = { | |
| "type": "object", | |
| "properties": { | |
| "shift_date_seen": { | |
| "type": "string", | |
| "nullable": True, | |
| "description": "A date written on the form header, as YYYY-MM-DD if you can determine it, else null.", | |
| }, | |
| "rows": { | |
| "type": "array", | |
| "description": "One entry per legible handwritten cell. Skip cells that are blank or fully illegible.", | |
| "items": { | |
| "type": "object", | |
| "properties": { | |
| "unit_name_seen": { | |
| "type": "string", | |
| "description": "The row's unit label exactly as written or printed on the form, verbatim.", | |
| }, | |
| "matched_unit_name": { | |
| "type": "string", | |
| "nullable": True, | |
| "description": "The single best match for this row from the provided list of known unit names, or null if none plausibly matches.", | |
| }, | |
| "slot_index": { | |
| "type": "integer", | |
| "nullable": True, | |
| "description": "0-11 index into the 12 time-slot columns (08:00,10:00,12:00,14:00,16:00,18:00,20:00,22:00,00:00,02:00,04:00,06:00), or null if the column is unclear.", | |
| }, | |
| "value_c": { | |
| "type": "number", | |
| "nullable": True, | |
| "description": "The handwritten numeric temperature in that cell, or null if it can't be read as a number.", | |
| }, | |
| "confidence": {"type": "number", "description": "0.0 to 1.0."}, | |
| "raw_text_seen": {"type": "string", "description": "The literal characters seen in the cell."}, | |
| }, | |
| "required": ["unit_name_seen", "slot_index", "value_c", "confidence", "raw_text_seen"], | |
| }, | |
| }, | |
| }, | |
| "required": ["rows"], | |
| } | |
| def build_form_extraction_prompt(unit_names: list) -> str: | |
| """Builds the prompt for bulk-extracting a photographed paper form. | |
| `unit_names` is the list of known active ctk_units names, given as context | |
| so the model can match handwritten row labels against real units instead | |
| of inventing names. | |
| """ | |
| names_block = "\n".join(f"- {n}" for n in unit_names) | |
| return f"""You are transcribing a photographed paper "ROOM TEMPERATURE RECORD" form from a | |
| meat-processing plant's cold-chain log. The form is a grid: each row is a chiller, | |
| freezer, or container, and there are 12 fixed time-slot columns across the top, in | |
| this exact order: 08:00, 10:00, 12:00, 14:00, 16:00, 18:00, 20:00, 22:00, 00:00, | |
| 02:00, 04:00, 06:00 (an 8am-to-8am shift). Entries are handwritten in pen or pencil. | |
| The plant's known unit names are: | |
| {names_block} | |
| For every cell that has a legible handwritten number in it, return one row with: | |
| the row's label exactly as written (unit_name_seen), your best match for that | |
| label against the list above (matched_unit_name, or null if nothing plausibly | |
| matches), which of the 12 columns it's in (slot_index, 0-11, matching the order | |
| above), and the numeric value (value_c). Skip cells that are blank or you truly | |
| cannot read. If the form's header has a date written on it, report your best | |
| reading of it as shift_date_seen (YYYY-MM-DD), else null. | |
| Handwriting varies a lot in legibility. Do your best on every cell rather than | |
| skipping anything you can make a reasonable guess at, but lower confidence | |
| accordingly rather than guessing wildly. A human reviews every row you return | |
| before anything is saved, so it's fine to be uncertain as long as it's reflected | |
| in confidence. | |
| Return your transcription as structured JSON matching the provided schema.""" | |
| def extract_form(image_bytes: bytes, mime_type: str, unit_names: list) -> dict: | |
| """Calls Gemini vision on a photographed paper form and returns the parsed | |
| {shift_date_seen, rows: [...]} dict. See FORM_SCHEMA for the row shape.""" | |
| client = get_client() | |
| prompt = build_form_extraction_prompt(unit_names) | |
| resp = client.models.generate_content( | |
| model=MODEL_NAME, | |
| contents=[ | |
| prompt, | |
| types.Part.from_bytes(data=image_bytes, mime_type=mime_type), | |
| ], | |
| config=types.GenerateContentConfig( | |
| response_mime_type="application/json", | |
| response_schema=FORM_SCHEMA, | |
| ), | |
| ) | |
| raw_json = "".join(part.text or "" for part in resp.candidates[0].content.parts) | |
| parsed = json.loads(raw_json) | |
| parsed["model"] = MODEL_NAME | |
| parsed["prompt_version"] = PROMPT_VERSION | |
| return parsed | |
| def extract_reading(image_bytes: bytes, mime_type: str, unit: dict) -> dict: | |
| """Calls Gemini vision on the photo and returns the parsed reading dict.""" | |
| client = get_client() | |
| prompt = build_extraction_prompt(unit) | |
| resp = client.models.generate_content( | |
| model=MODEL_NAME, | |
| contents=[ | |
| prompt, | |
| types.Part.from_bytes(data=image_bytes, mime_type=mime_type), | |
| ], | |
| config=types.GenerateContentConfig( | |
| response_mime_type="application/json", | |
| response_schema=READING_SCHEMA, | |
| ), | |
| ) | |
| raw_json = "".join(part.text or "" for part in resp.candidates[0].content.parts) | |
| parsed = json.loads(raw_json) | |
| parsed["model"] = MODEL_NAME | |
| parsed["prompt_version"] = PROMPT_VERSION | |
| return parsed | |