image_caption.py

from google.genai import types
from google import genai
import os
import json
import time
import random
from pydantic import BaseModel
from tqdm import tqdm


class MedicalFormat(BaseModel):
    hr: int
    sp02: int
    abp_systolic: int
    abp_diastolic: int
    abp_map: int
    etco2: int
    etco2_unit: str
    rr: int
    nibp_systolic: int
    nibp_diastolic: int
    nibp_map: int


client = genai.Client(api_key="")
image_folder = "images"
data_file = "data.json"


def image_caption(image_path, max_retries=5):
    if not os.path.exists(image_path):
        return None
        
    with open(image_path, "rb") as f:
        image_bytes = f.read()

    for attempt in range(max_retries):
        try:
            response = client.models.generate_content(
                model="gemini-2.5-flash",
                config={"response_mime_type": "application/json", "response_schema": MedicalFormat},
                contents=[
                    types.Part.from_bytes(
                        data=image_bytes,
                        mime_type="image/jpeg",
                    ),
                    "Extract all numerical patient vital signs from this medical monitor display as JSON.",
                ],
            )
            return response.text
        except Exception as e:
            wait = (2 ** attempt) + random.uniform(0, 1)
            print(f"Error: {e}. Retrying in {wait:.1f}s...")
            time.sleep(wait)

    print(f"❌ Skipping {image_path} after {max_retries} retries.")
    return None


def load_results():
    if not os.path.exists(data_file):
        return []
    try:
        with open(data_file, "r") as f:
            return json.load(f)
    except (json.JSONDecodeError, FileNotFoundError):
        return []


def save_results(results):
    with open(data_file, "w") as f:
        json.dump(results, f, indent=4)


def main():
    if not os.path.exists(image_folder):
        print(f"❌ Image folder '{image_folder}' not found.")
        return
        
    results = load_results()
    processed = {item["image_path"] for item in results}
    
    image_files = [f for f in os.listdir(image_folder)]
    
    for image in tqdm(image_files, desc="Processing images"):
        image_path = os.path.join(image_folder, image)
        
        if image_path in processed:
            print(f"⏩ Skipping {image} (already processed).")
            continue

        caption = image_caption(image_path)
        if caption is None:
            continue

        try:
            record = {"image_path": image_path, "data": json.loads(caption)}
            results.append(record)
            save_results(results)
            print(f"✅ Processed {image}")
        except json.JSONDecodeError as e:
            print(f"⚠️ Failed to parse JSON for {image}: {e}")
            continue


if __name__ == "__main__":
    main()