Spaces:
Runtime error
Runtime error
iamrobotbear
commited on
Commit
•
495ac6b
1
Parent(s):
4e52bb8
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,136 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
import torch
|
3 |
+
from PIL import Image
|
4 |
+
import pandas as pd
|
5 |
+
from lavis.models import load_model_and_preprocess
|
6 |
+
from lavis.processors import load_processor
|
7 |
+
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor
|
8 |
+
import tensorflow as tf
|
9 |
+
import tensorflow_hub as hub
|
10 |
+
import io
|
11 |
+
from sklearn.metrics.pairwise import cosine_similarity
|
12 |
+
import tempfile # Add this import
|
13 |
+
import logging
|
14 |
+
|
15 |
+
# Configure logging
|
16 |
+
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
17 |
+
|
18 |
+
# Load model and preprocessors for Image-Text Matching (LAVIS)
|
19 |
+
device = torch.device("cuda") if torch.cuda.is_available() else "cpu"
|
20 |
+
model_itm, vis_processors, text_processors = load_model_and_preprocess("blip2_image_text_matching", "pretrain", device=device, is_eval=True)
|
21 |
+
|
22 |
+
# Load tokenizer and model for Image Captioning (TextCaps)
|
23 |
+
git_processor_large_textcaps = AutoProcessor.from_pretrained("microsoft/git-large-r-textcaps")
|
24 |
+
git_model_large_textcaps = AutoModelForCausalLM.from_pretrained("microsoft/git-large-r-textcaps")
|
25 |
+
|
26 |
+
# Load Universal Sentence Encoder model for textual similarity calculation
|
27 |
+
embed = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
|
28 |
+
|
29 |
+
# Define a function to compute textual similarity between caption and statement
|
30 |
+
def compute_textual_similarity(caption, statement):
|
31 |
+
# Convert caption and statement into sentence embeddings
|
32 |
+
caption_embedding = embed([caption])[0].numpy()
|
33 |
+
statement_embedding = embed([statement])[0].numpy()
|
34 |
+
|
35 |
+
# Calculate cosine similarity between sentence embeddings
|
36 |
+
similarity_score = cosine_similarity([caption_embedding], [statement_embedding])[0][0]
|
37 |
+
return similarity_score
|
38 |
+
|
39 |
+
# Read statements from the external file 'statements.txt'
|
40 |
+
with open('statements.txt', 'r') as file:
|
41 |
+
statements = file.read().splitlines()
|
42 |
+
|
43 |
+
# Function to compute ITM scores for the image-statement pair
|
44 |
+
def compute_itm_score(image, statement):
|
45 |
+
logging.info('Starting compute_itm_score')
|
46 |
+
pil_image = Image.fromarray(image.astype('uint8'), 'RGB')
|
47 |
+
img = vis_processors["eval"](pil_image.convert("RGB")).unsqueeze(0).to(device)
|
48 |
+
# Pass the statement text directly to model_itm
|
49 |
+
itm_output = model_itm({"image": img, "text_input": statement}, match_head="itm")
|
50 |
+
itm_scores = torch.nn.functional.softmax(itm_output, dim=1)
|
51 |
+
score = itm_scores[:, 1].item()
|
52 |
+
logging.info('Finished compute_itm_score')
|
53 |
+
return score
|
54 |
+
|
55 |
+
def generate_caption(processor, model, image):
|
56 |
+
logging.info('Starting generate_caption')
|
57 |
+
inputs = processor(images=image, return_tensors="pt").to(device)
|
58 |
+
generated_ids = model.generate(pixel_values=inputs.pixel_values, max_length=50)
|
59 |
+
generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
|
60 |
+
logging.info('Finished generate_caption')
|
61 |
+
return generated_caption
|
62 |
+
|
63 |
+
def save_dataframe_to_csv(df):
|
64 |
+
csv_buffer = io.StringIO()
|
65 |
+
df.to_csv(csv_buffer, index=False)
|
66 |
+
csv_string = csv_buffer.getvalue()
|
67 |
+
|
68 |
+
# Save the CSV string to a temporary file
|
69 |
+
with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".csv") as temp_file:
|
70 |
+
temp_file.write(csv_string)
|
71 |
+
temp_file_path = temp_file.name # Get the file path
|
72 |
+
|
73 |
+
# Return the file path (no need to reopen the file with "rb" mode)
|
74 |
+
return temp_file_path
|
75 |
+
|
76 |
+
# Main function to perform image captioning and image-text matching
|
77 |
+
def process_images_and_statements(image):
|
78 |
+
logging.info('Starting process_images_and_statements')
|
79 |
+
|
80 |
+
# Generate image caption for the uploaded image using git-large-r-textcaps
|
81 |
+
caption = generate_caption(git_processor_large_textcaps, git_model_large_textcaps, image)
|
82 |
+
|
83 |
+
# Define weights for combining textual similarity score and image-statement ITM score (adjust as needed)
|
84 |
+
weight_textual_similarity = 0.5
|
85 |
+
weight_statement = 0.5
|
86 |
+
|
87 |
+
# Initialize an empty list to store the results
|
88 |
+
results_list = []
|
89 |
+
|
90 |
+
# Loop through each predefined statement
|
91 |
+
for statement in statements:
|
92 |
+
# Compute textual similarity between caption and statement
|
93 |
+
textual_similarity_score = (compute_textual_similarity(caption, statement) * 100) # Multiply by 100
|
94 |
+
|
95 |
+
# Compute ITM score for the image-statement pair
|
96 |
+
itm_score_statement = (compute_itm_score(image, statement) * 100) # Multiply by 100
|
97 |
+
|
98 |
+
# Combine the two scores using a weighted average
|
99 |
+
final_score = ((weight_textual_similarity * textual_similarity_score) +
|
100 |
+
(weight_statement * itm_score_statement))
|
101 |
+
|
102 |
+
# Append the result to the results_list
|
103 |
+
results_list.append({
|
104 |
+
'Statement': statement,
|
105 |
+
'Generated Caption': caption, # Include the generated caption
|
106 |
+
'Textual Similarity Score': f"{textual_similarity_score:.2f}%", # Format as percentage with two decimal places
|
107 |
+
'ITM Score': f"{itm_score_statement:.2f}%", # Format as percentage with two decimal places
|
108 |
+
'Final Combined Score': f"{final_score:.2f}%" # Format as percentage with two decimal places
|
109 |
+
})
|
110 |
+
|
111 |
+
# Convert the results_list to a DataFrame using pandas.concat
|
112 |
+
results_df = pd.concat([pd.DataFrame([result]) for result in results_list], ignore_index=True)
|
113 |
+
|
114 |
+
logging.info('Finished process_images_and_statements')
|
115 |
+
|
116 |
+
# Save results_df to a CSV file
|
117 |
+
csv_results = save_dataframe_to_csv(results_df)
|
118 |
+
|
119 |
+
# Return both the DataFrame and the CSV data for the Gradio interface
|
120 |
+
return results_df, csv_results # <--- Return results_df and csv_results
|
121 |
+
|
122 |
+
# Gradio interface
|
123 |
+
image_input = gr.inputs.Image()
|
124 |
+
output_df = gr.outputs.Dataframe(type="pandas", label="Results")
|
125 |
+
output_csv = gr.outputs.File(label="Download CSV")
|
126 |
+
|
127 |
+
iface = gr.Interface(
|
128 |
+
fn=process_images_and_statements,
|
129 |
+
inputs=image_input,
|
130 |
+
outputs=[output_df, output_csv], # Include both the DataFrame and CSV file outputs
|
131 |
+
title="Image Captioning and Image-Text Matching",
|
132 |
+
theme='sudeepshouche/minimalist',
|
133 |
+
css=".output { flex-direction: column; } .output .outputs { width: 100%; }" # Custom CSS
|
134 |
+
)
|
135 |
+
|
136 |
+
iface.launch()
|