|
import gradio as gr |
|
import torch |
|
from PIL import Image |
|
from transformers import MllamaForConditionalGeneration, AutoProcessor |
|
|
|
|
|
model_id = "0llheaven/Llama-3.2-11B-Vision-Radiology-mini" |
|
|
|
model = MllamaForConditionalGeneration.from_pretrained( |
|
model_id, |
|
torch_dtype=torch.bfloat16, |
|
device_map="auto", |
|
) |
|
processor = AutoProcessor.from_pretrained(model_id) |
|
|
|
|
|
def generate_caption(image): |
|
|
|
image = image.convert("RGB") |
|
|
|
instruction = "You are an expert radiographer. Describe accurately what you see in this image." |
|
messages = [ |
|
{"role": "user", "content": [ |
|
{"type": "image"}, |
|
{"type": "text", "text": instruction} |
|
]} |
|
] |
|
|
|
input_text = processor.apply_chat_template(messages, add_generation_prompt=True) |
|
inputs = processor( |
|
image, |
|
input_text, |
|
add_special_tokens=False, |
|
return_tensors="pt" |
|
).to(model.device) |
|
|
|
|
|
output = model.generate(**inputs, max_new_tokens=256, use_cache=True, temperature=1.5, min_p=0.1) |
|
return processor.decode(output[0]) |
|
|
|
|
|
with gr.Blocks() as demo: |
|
gr.Markdown("# Radiology Image Captioning") |
|
with gr.Row(): |
|
image_input = gr.Image(type="pil", label="Upload Image") |
|
output_text = gr.Textbox(label="Generated Caption") |
|
title="Medical Vision Analysis" |
|
generate_button = gr.Button("Generate Caption") |
|
|
|
|
|
generate_button.click(fn=generate_caption, inputs=image_input, outputs=output_text) |
|
|
|
|
|
demo.launch() |
|
|
