medgemma-cataract-surgical-analysis

from transformers import AutoProcessor, AutoModelForImageTextToText
from PIL import Image
import torch

# Load the merged model and processor
model_path = "mehti/medgemma-cataract-surgical-analysis"
processor = AutoProcessor.from_pretrained(model_path)
model = AutoModelForImageTextToText.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
model.eval()

# Load a surgical frame (replace with your image)
image = Image.open("surgical_frame.jpg").convert("RGB")

messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "image": image},
            {
                "type": "text",
                "text": (
                    "You are an expert ophthalmic surgeon reviewing a frame from a cataract surgery video. "
                    "Analyze the surgical scene and provide a chain-of-thought reasoning followed by a "
                    "clear instruction for a surgical resident.\n\n"
                    "Format your response as:\nThinking Process:\n<your reasoning>\n\nFinal Answer:\n<your instruction>"
                ),
            },
        ],
    }
]

inputs = processor.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
).to(model.device, dtype=torch.bfloat16)

with torch.inference_mode():
    output = model.generate(**inputs, max_new_tokens=512, do_sample=False)

decoded = processor.decode(output[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
print(decoded)