iamrobotbear commited on
Commit
495ac6b
1 Parent(s): 4e52bb8

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +136 -0
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()