Update app.py

app.py

@@ -51,7 +51,6 @@ def compute_itm_score(image, statement):
     score = itm_scores[:, 1].item()
     logging.info('Finished compute_itm_score')
     return score
-
 def generate_caption(processor, model, image):
     logging.info('Starting generate_caption')
     inputs = processor(images=image, return_tensors="pt").to(device)
@@ -74,7 +73,7 @@ def save_dataframe_to_csv(df):
     return temp_file_path
 
 # Main function to perform image captioning and image-text matching
-def process_images_and_statements(image):
+def process_images_and_statements(image, filename):
     logging.info('Starting process_images_and_statements')
 
     # Generate image caption for the uploaded image using git-large-r-textcaps
@@ -99,10 +98,11 @@ def process_images_and_statements(image):
         final_score = ((weight_textual_similarity * textual_similarity_score) +
                        (weight_statement * itm_score_statement))
 
-        # Append the result to the results_list
+        # Append the result to the results_list, including the image filename
         results_list.append({
+            'Image Filename': filename,  # Add the image filename to the output
             'Statement': statement,
-            'Generated Caption': caption, # Include the generated caption
+            'Generated Caption': caption,
             'Textual Similarity Score': f"{textual_similarity_score:.2f}%", # Format as percentage with two decimal places
             'ITM Score': f"{itm_score_statement:.2f}%", # Format as percentage with two decimal places
             'Final Combined Score': f"{final_score:.2f}%" # Format as percentage with two decimal places
@@ -119,18 +119,24 @@ def process_images_and_statements(image):
     # Return both the DataFrame and the CSV data for the Gradio interface
     return results_df, csv_results # <--- Return results_df and csv_results
 
+# Define a function to handle the user's input
+def handle_input(image, filename):
+    # Call the process_images_and_statements function with the image and filename
+    return process_images_and_statements(image, filename.name)
 # Gradio interface
 image_input = gr.inputs.Image()
+file_input = gr.inputs.File()  # Add an input element to allow users to provide the image filename
 output_df = gr.outputs.Dataframe(type="pandas", label="Results")
 output_csv = gr.outputs.File(label="Download CSV")
 
 iface = gr.Interface(
-    fn=process_images_and_statements,
-    inputs=image_input,
-    outputs=[output_df, output_csv], # Include both the DataFrame and CSV file outputs
+    fn=handle_input,  # Use the handle_input function to handle user input
+    inputs=[image_input, file_input],  # Include both the image and filename inputs
+    outputs=[output_df, output_csv],
     title="Image Captioning and Image-Text Matching",
     theme='sudeepshouche/minimalist',
     css=".output { flex-direction: column; } .output .outputs { width: 100%; }" # Custom CSS
 )
 
-iface.launch()
+# Launch the Gradio interface
+iface.launch()