app.py

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



# Caption generation functions
def generate_caption_clipgpt(image, max_tokens, temperature):
  caption = clipGPT.generate_caption_clipgpt(image, max_tokens, temperature)
  return caption

def generate_caption_vitgpt(image, max_tokens, temperature):
  caption = vitGPT.generate_caption(image, max_tokens, temperature)
  return caption

def generate_caption_vitCoAtt(image):
  caption = ViTCoAtt.CaptionSampler.main(image)
  return caption


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():
         
    image = gr.Image(label="Upload Chest X-ray", type="pil")   
         
           
    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(): # Column for dropdowns and model choice
        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, max_tokens, temperature):
    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."   



examples = [f"example{i}.png" for i in range(1,3)]

interface = gr.Interface(
        fn=predict,
        inputs = [image, model_choice, max_tokens, temperature],
        theme="gradio/monochrome",
        outputs=caption,
        examples = examples
    )

interface.launch(debug=True)