Update app.py

app.py

@@ -38,46 +38,47 @@ with gr.Row():
 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) 
+        temperature = gr.Slider(0.5, 0.9, step=0.1, label="Temperature", value=0.9)
+        imgID = gr.Dropdown(list(range(1,6)), label="Example Image Selected", value=1)
 
     model_choice = gr.Radio(["CLIP-GPT2", "ViT-GPT2", "ViT-CoAttention", "Baseline Model CNN-RNN"], label="Select Model") 
     generate_button = gr.Button("Generate Caption") 
     
 
-real_captions = {"0" : "No acute cardiopulmonary abnormality. Low lung volumes. Heart size and mediastinal contour within normal limits. No focal air space consolidation, pneumothorax, or pleural effusion. Mild thoracic spine degenerative change.", 
-                 "1":"Left basilar atelectasis and/or infiltrate, with no radiographic evidence of tuberculosis. Heart size upper limits of normal. Small amount of left basilar airspace disease. The right lung is clear. There are no cavitary lesions seen. No pneumothorax. No pleural effusions",
-                 "2":"Cardiomegaly and small bilateral pleural effusions. Abnormal pulmonary opacities most suggestive of pulmonary edema, primary differential diagnosis includes infection and aspiration, clinical correlation recommended Moderate-to-marked enlargement of the cardiac silhouette, mediastinal contours appear similar to prior. Mild bilateral posterior sulcus blunting, interstitial and alveolar opacities greatest in the central lungs and bases with indistinct vascular margination.",
-                 "3":"Severe cardiomegaly. Limited mediastinal evaluation given body habitus and lordotic projection. Recommend XXXX for further evaluation of mediastinum given T/Spine injury noted on C/Spine imaging. Critical result notification documented through Primordial. Lordotic projection and large body habitus. Limited mediastinal evaluation. Severe cardiomegaly. No visualized pneumothorax. No large effusion or airspace disease. No fracture."}
+real_captions = {"1" : "No acute cardiopulmonary abnormality. Low lung volumes. Heart size and mediastinal contour within normal limits. No focal air space consolidation, pneumothorax, or pleural effusion. Mild thoracic spine degenerative change.", 
+                 "2":"Left basilar atelectasis and/or infiltrate, with no radiographic evidence of tuberculosis. Heart size upper limits of normal. Small amount of left basilar airspace disease. The right lung is clear. There are no cavitary lesions seen. No pneumothorax. No pleural effusions",
+                 "3":"Cardiomegaly and small bilateral pleural effusions. Abnormal pulmonary opacities most suggestive of pulmonary edema, primary differential diagnosis includes infection and aspiration, clinical correlation recommended Moderate-to-marked enlargement of the cardiac silhouette, mediastinal contours appear similar to prior. Mild bilateral posterior sulcus blunting, interstitial and alveolar opacities greatest in the central lungs and bases with indistinct vascular margination.",
+                 "4":"Severe cardiomegaly. Limited mediastinal evaluation given body habitus and lordotic projection. Recommend XXXX for further evaluation of mediastinum given T/Spine injury noted on C/Spine imaging. Critical result notification documented through Primordial. Lordotic projection and large body habitus. Limited mediastinal evaluation. Severe cardiomegaly. No visualized pneumothorax. No large effusion or airspace disease. No fracture."}
 
-imgIDs = {"0":"/content/drive/MyDrive/cnn-rnn/NLMCXR_png/CXR192_IM-0598_0",
-          "1":"/content/drive/MyDrive/cnn-rnn/NLMCXR_png/CXR194_IM-0609_0",
-          "2":"/content/drive/MyDrive/cnn-rnn/NLMCXR_png/CXR2637_IM-1122_0",
-          "3":"/content/drive/MyDrive/cnn-rnn/NLMCXR_png/CXR1111_IM-0077_0"}
+imgIDs = {"1":"/content/drive/MyDrive/cnn-rnn/NLMCXR_png/CXR192_IM-0598_0",
+          "2":"/content/drive/MyDrive/cnn-rnn/NLMCXR_png/CXR194_IM-0609_0",
+          "3":"/content/drive/MyDrive/cnn-rnn/NLMCXR_png/CXR2637_IM-1122_0",
+          "4":"/content/drive/MyDrive/cnn-rnn/NLMCXR_png/CXR1111_IM-0077_0"}
 
 
 caption = gr.Textbox(label="Generated Caption")
 real_caption = gr.Textbox(label="Actual Caption")
 
-def predict(img, model_name, max_tokens, temperature, examples):
+def predict(img, model_name, max_tokens, temperature, imgID):
     if model_name == "CLIP-GPT2":
-        return generate_caption_clipgpt(img, max_tokens, temperature), getCaption(examples)
+        return generate_caption_clipgpt(img, max_tokens, temperature), getCaption(imgID)
     elif model_name == "ViT-GPT2":
-        return generate_caption_vitgpt(img, max_tokens, temperature), getCaption(examples)
+        return generate_caption_vitgpt(img, max_tokens, temperature), getCaption(imgID)
     elif model_name == "ViT-CoAttention":
-        return generate_caption_vitCoAtt(img),  getCaption(examples)
+        return generate_caption_vitCoAtt(img),  getCaption(imgID)
     elif model_name == "Baseline Model CNN-RNN":
-        img = getImageID(examples)
-        return generate_caption_cnnrnn(img), getCaption(examples)
+        img = getImageID(imgID)
+        return generate_caption_cnnrnn(img), getCaption(imgID)
     else:
         return "Caption generation for this model is not yet implemented."   
 
-def getCaption(examples):
-    print(real_captions[examples[1]])
-    return real_captions[examples[1]]
+def getCaption(imgID):
+    print(real_captions[imgID])
+    return real_captions[imgID]
     
-def getImageID(examples):
-    print(imgIDs[examples[1]])
-    return imgIDs[examples[1]]
+def getImageID(imgID):
+    print(imgIDs[imgID])
+    return imgIDs[imgID]
 
 examples = [[f"example{i}.jpg"] for i in range(1,7)]
 
@@ -86,7 +87,7 @@ title = "MedViT: A Vision Transformer-Driven Method for Generating Medical Repor
 
 interface = gr.Interface(
         fn=predict,
-        inputs = [image, model_choice, max_tokens, temperature, examples],
+        inputs = [image, model_choice, max_tokens, temperature, imgID],
         theme="sudeepshouche/minimalist",
         outputs=[caption,real_caption],
         examples = examples,