File size: 2,523 Bytes
830c7d1 | 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 | import gradio as gr
import torch
from PIL import Image
import os
from transformers import AutoProcessor, LlavaNextForConditionalGeneration
# Model config
MODEL_ID = "OralGPT/OralGPT-Omni-7B-Instruct" # Based on README news
def load_model():
print(f"Loading model {MODEL_ID}...")
processor = AutoProcessor.from_pretrained(MODEL_ID)
model = LlavaNextForConditionalGeneration.from_pretrained(
MODEL_ID,
torch_dtype=torch.float16,
low_cpu_mem_usage=True,
device_map="auto"
)
return processor, model
# Global variables for the model and processor
processor = None
model = None
def predict(image, text):
global processor, model
if processor is None or model is None:
processor, model = load_model()
if image is None:
return "Please upload a dental image."
# Prepare prompt following LLaVA-v1.6 / LLaVA-Next style (which OralGPT-Omni uses)
# Note: OralGPT-Omni is based on Qwen/Vicuna usually.
# The prompt template might need adjustment based on the specific base model.
prompt = f"A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions. USER: <image>\n{text} ASSISTANT:"
inputs = processor(text=prompt, images=image, return_tensors="pt").to(model.device, torch.float16)
with torch.no_grad():
output = model.generate(**inputs, max_new_tokens=512)
response = processor.decode(output[0], skip_special_tokens=True)
# Post-process to extract only the assistant's part
if "ASSISTANT:" in response:
response = response.split("ASSISTANT:")[-1].strip()
return response
# Gradio Interface
description = """
# 🦷 OralGPT-Omni: Digital Dentistry Assistant
Upload a dental image (X-ray, photograph, etc.) and ask a question.
OralGPT is specialized in dental multimodal analysis including Panoramic X-rays, Periapical radiographs, and more.
"""
demo = gr.Interface(
fn=predict,
inputs=[
gr.Image(type="pil", label="Dental Image"),
gr.Textbox(label="Question", placeholder="Describe this dental X-ray...")
],
outputs=gr.Textbox(label="OralGPT Response"),
title="OralGPT-Omni Demo",
description=description,
examples=[
["OralGPT/assets/mmoral-logo.png", "What is shown in this logo?"]
]
)
if __name__ == "__main__":
demo.launch()
|