|
import gradio as gr |
|
from PIL import Image |
|
import clipGPT |
|
import vitGPT |
|
import skimage.io as io |
|
import PIL.Image |
|
import difflib |
|
import ViTCoAtt |
|
from build_vocab import Vocabulary |
|
|
|
|
|
|
|
|
|
def generate_caption_clipgpt(image, max_tokens, temperature): |
|
caption = clipGPT.generate_caption_clipgpt(image) |
|
return caption |
|
|
|
def generate_caption_vitgpt(image, max_tokens, temperature): |
|
caption = vitGPT.generate_caption(image) |
|
return caption |
|
|
|
def generate_caption_vitCoAtt(image): |
|
caption = ViTCoAtt.CaptionSampler.main(image) |
|
return caption |
|
|
|
|
|
with gr.Blocks() as demo: |
|
|
|
gr.HTML("<h1 style='text-align: center;'>MedViT: A Vision Transformer-Driven Method for Generating Medical Reports 🏥🤖</h1>") |
|
gr.HTML("<p style='text-align: center;'>You can generate captions by uploading an X-Ray and selecting a model of your choice below</p>") |
|
|
|
|
|
with gr.Row(): |
|
sample_images = [ |
|
'https://imgur.com/W1pIr9b', |
|
'https://imgur.com/MLJaWnf', |
|
'https://imgur.com/6XymFW1', |
|
'https://imgur.com/zdPjZZ1', |
|
'https://imgur.com/DKUlZbF' |
|
] |
|
|
|
|
|
image = gr.Image(label="Upload Chest X-ray", type="pil") |
|
|
|
sample_images_gallery = gr.Gallery(value = sample_images,label="Sample Images") |
|
|
|
gr.HTML("<p style='text-align: center;'> Please select the Number of Max Tokens and Temperature setting, if you are testing CLIP GPT2 and VIT GPT2 Models</p>") |
|
|
|
|
|
with gr.Row(): |
|
|
|
with gr.Column(): |
|
max_tokens = gr.Dropdown(list(range(50, 101)), label="Max Tokens", value=75) |
|
temperature = gr.Slider(0.5, 0.9, step=0.1, label="Temperature", value=0.7) |
|
|
|
model_choice = gr.Radio(["CLIP-GPT2", "ViT-GPT2", "ViT-CoAttention"], label="Select Model") |
|
|
|
generate_button = gr.Button("Generate Caption") |
|
|
|
|
|
|
|
caption = gr.Textbox(label="Generated Caption") |
|
|
|
def predict(img, model_name): |
|
if model_name == "CLIP-GPT2": |
|
return generate_caption_clipgpt(img, max_tokens, temperature) |
|
elif model_name == "ViT-GPT2": |
|
return generate_caption_vitgpt(img, max_tokens, temperature) |
|
elif model_name == "ViT-CoAttention": |
|
return generate_caption_vitCoAtt(img) |
|
else: |
|
return "Caption generation for this model is not yet implemented." |
|
|
|
|
|
|
|
generate_button.click(predict, [image, model_choice, max_tokens, temperature], caption) |
|
sample_images_gallery.change(predict, [sample_images_gallery, model_choice, max_tokens, temperature], caption) |
|
|
|
|
|
demo.launch() |
|
|
