Update app.py

app.py

@@ -1,22 +1,19 @@
 import gradio as gr
 from PIL import Image
 import clipGPT
+import vitGPT
 import skimage.io as io
 import PIL.Image
 
-# Define model loading functions (if needed)
-def load_model_1():  # CLIP-GPT2
-    # Load model components here if necessary
-    return None 
-
-# ... load_model_2(), load_model_3() - Define if and when needed
 
 # Caption generation functions
 def generate_caption_clipgpt(image):
     caption = clipGPT.generate_caption_clipgpt(image)
     return caption
 
-# ... Add more caption generation functions for future models
+def generate_caption_vitgpt(image):
+    caption = vitGPT.generate_caption(image)
+    return caption
 
 # Sample image paths
 sample_images = [
@@ -27,15 +24,12 @@ sample_images = [
     "CXR195_IM-0618-1001.png"
 ]
 
-def load_image(path):
-    return Image.open(path)
-
 # Gradio interface 
 with gr.Blocks() as demo:
     with gr.Row():
         image = gr.Image(label="Upload Chest X-ray") 
         sample_images_gallery = gr.Gallery(
-            [load_image(path) for path in sample_images], 
+            sample_images, 
             label="Sample Images",
         )  
     with gr.Row():  
@@ -45,9 +39,9 @@ with gr.Blocks() as demo:
 
     def predict(img, model_name):
         if model_name == "CLIP-GPT2":
-            print(img)
             return generate_caption_clipgpt(img)
-        # Add elif blocks for "ViT-GPT2", "ViT-CoAttention"  as you implement them
+        elif model_name == "ViT-GPT2":
+            return generate_caption_vitgpt(img)
         else:
             return "Caption generation for this model is not yet implemented."