going back to pandas.concat

app.py

@@ -1,92 +1,136 @@
 import gradio as gr
+import torch
+from PIL import Image
+import pandas as pd
+from lavis.models import load_model_and_preprocess
+from lavis.processors import load_processor
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor
 import tensorflow as tf
 import tensorflow_hub as hub
-import numpy as np
-import pandas as pd
-from transformers import GitProcessor, GitModel, GitConfig
-from PIL import Image
-import logging 
+import io
+from sklearn.metrics.pairwise import cosine_similarity
+import tempfile # Add this import
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+# Load model and preprocessors for Image-Text Matching (LAVIS)
+device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
+model_itm, vis_processors, text_processors = load_model_and_preprocess("blip2_image_text_matching", "pretrain", device=device, is_eval=True)
+
+# Load tokenizer and model for Image Captioning (TextCaps)
+git_processor_large_textcaps = AutoProcessor.from_pretrained("microsoft/git-large-r-textcaps")
+git_model_large_textcaps = AutoModelForCausalLM.from_pretrained("microsoft/git-large-r-textcaps")
 
-# Load models and processors
-git_config = GitConfig.from_pretrained("microsoft/git-large-r")
-git_processor_large_textcaps = GitProcessor.from_pretrained("microsoft/git-large-r")
-git_model_large_textcaps = GitModel.from_pretrained("microsoft/git-large-r")
-itm_model = hub.load("https://tfhub.dev/google/LaViT/1")
-use_model = hub.load("https://tfhub.dev/google/universal-sentence-encoder-large/5")
+# Load Universal Sentence Encoder model for textual similarity calculation
+embed = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
+
+# Define a function to compute textual similarity between caption and statement
+def compute_textual_similarity(caption, statement):
+    # Convert caption and statement into sentence embeddings
+    caption_embedding = embed([caption])[0].numpy()
+    statement_embedding = embed([statement])[0].numpy()
+
+    # Calculate cosine similarity between sentence embeddings
+    similarity_score = cosine_similarity([caption_embedding], [statement_embedding])[0][0]
+    return similarity_score
 
 # Read statements from the external file 'statements.txt'
 with open('statements.txt', 'r') as file:
     statements = file.read().splitlines()
 
-# Function to generate image caption
+# Function to compute ITM scores for the image-statement pair
+def compute_itm_score(image, statement):
+    logging.info('Starting compute_itm_score')
+    pil_image = Image.fromarray(image.astype('uint8'), 'RGB')
+    img = vis_processors["eval"](pil_image.convert("RGB")).unsqueeze(0).to(device)
+    # Pass the statement text directly to model_itm
+    itm_output = model_itm({"image": img, "text_input": statement}, match_head="itm")
+    itm_scores = torch.nn.functional.softmax(itm_output, dim=1)
+    score = itm_scores[:, 1].item()
+    logging.info('Finished compute_itm_score')
+    return score
+
 def generate_caption(processor, model, image):
-    inputs = processor(images=image, return_tensors="pt")
-    outputs = model(**inputs)
-    caption = processor.batch_decode(outputs.logits.argmax(-1), skip_special_tokens=True)
-    return caption[0]
+    logging.info('Starting generate_caption')
+    inputs = processor(images=image, return_tensors="pt").to(device)
+    generated_ids = model.generate(pixel_values=inputs.pixel_values, max_length=50)
+    generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    logging.info('Finished generate_caption')
+    return generated_caption
 
-# Function to compute textual similarity
-def compute_textual_similarity(caption, statement):
-    captions_embeddings = use_model([caption])[0].numpy()
-    statements_embeddings = use_model([statement])[0].numpy()
-    similarity_score = np.inner(captions_embeddings, statements_embeddings)
-    return similarity_score[0]
+def save_dataframe_to_csv(df):
+    csv_buffer = io.StringIO()
+    df.to_csv(csv_buffer, index=False)
+    csv_string = csv_buffer.getvalue()
 
-# Function to compute ITM score
-def compute_itm_score(image, statement):
-    image_features = itm_model(image)
-    statement_features = use_model([statement])[0].numpy()
-    similarity_score = np.inner(image_features, statement_features)
-    return similarity_score[0][0]
+    # Save the CSV string to a temporary file
+    with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".csv") as temp_file:
+        temp_file.write(csv_string)
+        temp_file_path = temp_file.name # Get the file path
 
-# Function to save DataFrame to CSV
-def save_dataframe_to_csv(df):
-    csv_data = df.to_csv(index=False)
-    return csv_data
-
-# Main function to perform image captioning and image-text matching for multiple images
-def process_images_and_statements(files):
-    all_results_list = []
-    
-    # If 'files' is a list, convert it to a dictionary
-    if isinstance(files, list):
-        files = {f.name: f for f in files}
-    
-    for file_name, image_file in files.items():
-        # Convert the image file to a PIL image
-        image = Image.open(image_file)
-        
-        caption = generate_caption(git_processor_large_textcaps, git_model_large_textcaps, image)
-        for statement in statements:
-            textual_similarity_score = compute_textual_similarity(caption, statement) * 100
-            itm_score_statement = compute_itm_score(image, statement) * 100
-            final_score = 0.5 * textual_similarity_score + 0.5 * itm_score_statement
-            all_results_list.append({
-                'Image File Name': file_name,  # Include the image file name
-                'Statement': statement,
-                'Generated Caption': caption,
-                'Textual Similarity Score': f"{textual_similarity_score:.2f}%",
-                'ITM Score': f"{itm_score_statement:.2f}%",
-                'Final Combined Score': f"{final_score:.2f}%"
-            })
-    results_df = pd.concat([pd.DataFrame([result]) for result in all_results_list], ignore_index=True)
+    # Return the file path (no need to reopen the file with "rb" mode)
+    return temp_file_path
+
+# Main function to perform image captioning and image-text matching
+def process_images_and_statements(image):
+    logging.info('Starting process_images_and_statements')
+
+    # Generate image caption for the uploaded image using git-large-r-textcaps
+    caption = generate_caption(git_processor_large_textcaps, git_model_large_textcaps, image)
+
+    # Define weights for combining textual similarity score and image-statement ITM score (adjust as needed)
+    weight_textual_similarity = 0.5
+    weight_statement = 0.5
+
+    # Initialize an empty list to store the results
+    results_list = []
+
+    # Loop through each predefined statement
+    for statement in statements:
+        # Compute textual similarity between caption and statement
+        textual_similarity_score = (compute_textual_similarity(caption, statement) * 100) # Multiply by 100
+
+        # Compute ITM score for the image-statement pair
+        itm_score_statement = (compute_itm_score(image, statement) * 100) # Multiply by 100
+
+        # Combine the two scores using a weighted average
+        final_score = ((weight_textual_similarity * textual_similarity_score) +
+                       (weight_statement * itm_score_statement))
+
+        # Append the result to the results_list
+        results_list.append({
+            'Statement': statement,
+            'Generated Caption': caption, # Include the generated 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
+        })
+
+    # Convert the results_list to a DataFrame using pandas.concat
+    results_df = pd.concat([pd.DataFrame([result]) for result in results_list], ignore_index=True)
+
+    logging.info('Finished process_images_and_statements')
+
+    # Save results_df to a CSV file
     csv_results = save_dataframe_to_csv(results_df)
-    return results_df, csv_results
 
-# Gradio interface with File input to receive
-# Gradio interface with File input to receive multiple images and file names
-image_input = gr.inputs.File(file_count="multiple", type="file", label="Upload Images")
+    # Return both the DataFrame and the CSV data for the Gradio interface
+    return results_df, csv_results # <--- Return results_df and csv_results
+
+# Gradio interface
+image_input = gr.inputs.Image()
 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],
+    outputs=[output_df, output_csv], # Include both the DataFrame and CSV file outputs
     title="Image Captioning and Image-Text Matching",
     theme='sudeepshouche/minimalist',
-    css=".output { flex-direction: column; } .output .outputs { width: 100%; }", # Custom CSS
-    capture_session=True,  # Capture errors and exceptions in Gradio interface
+    css=".output { flex-direction: column; } .output .outputs { width: 100%; }" # Custom CSS
 )
 
-iface.launch(debug=True)
+iface.launch()