Update app.py

app.py

@@ -11,26 +11,25 @@ from build_vocab import Vocabulary
 
 
 # Caption generation functions
-def generate_caption_clipgpt(image):
-    caption = clipGPT.generate_caption_clipgpt(image)
-    return caption
+def generate_caption_clipgpt(image, max_tokens, temperature):
+  caption = clipGPT.generate_caption_clipgpt(image, max_tokens, temperature)
+  return caption
 
-def generate_caption_vitgpt(image):
-    caption = vitGPT.generate_caption(image)
-    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
+  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',
@@ -38,36 +37,43 @@ with gr.Blocks() as demo:
         'https://imgur.com/6XymFW1',
         'https://imgur.com/zdPjZZ1',
         'https://imgur.com/DKUlZbF'
-            ]
-
-        
-        image = gr.Image(label="Upload Chest X-ray", type="pil")    
+          ]
         
+         
+        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():
-        model_choice = gr.Radio(["CLIP-GPT2", "ViT-GPT2", "ViT-CoAttention"], label="Select Model")
-
+     
+        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):
+    
+    def predict(img, model_name, max_tokens, temperature):
         if model_name == "CLIP-GPT2":
-            return generate_caption_clipgpt(img)
+          return generate_caption_clipgpt(img, max_tokens, temperature)
         elif model_name == "ViT-GPT2":
-            return generate_caption_vitgpt(img)
+          return generate_caption_vitgpt(img, max_tokens, temperature)
         elif model_name == "ViT-CoAttention":
-            return generate_caption_vitCoAtt(img)
+          return generate_caption_vitCoAtt(img)
         else:
-            return "Caption generation for this model is not yet implemented."     
-
-
-    # Event handlers
-    generate_button.click(predict, [image, model_choice], caption)  # Trigger prediction on button click 
-    sample_images_gallery.change(predict, [sample_images_gallery, model_choice], caption)  # Handle sample images
-
+          return "Caption generation for this model is not yet implemented."   
+        
+    
+    # Event handlers 
+    generate_button.click(predict, [image, model_choice, max_tokens, temperature], caption) 
+    sample_images_gallery.click(predict, [sample_images_gallery, model_choice, max_tokens, temperature], caption)  
+    
 
-demo.launch() 
+demo.launch()