Spaces:
Build error
Build error
| 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 io | |
| import os | |
| import numpy as np | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| import tempfile | |
| import shutil | |
| 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 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 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): | |
| 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 | |
| def save_dataframe_to_csv(df): | |
| csv_buffer = io.StringIO() | |
| df.to_csv(csv_buffer, index=False) | |
| csv_string = csv_buffer.getvalue() | |
| # 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 | |
| # Return the file path (no need to reopen the file with "rb" mode) | |
| return temp_file_path | |
| def process_images_and_statements(image_file): | |
| with tempfile.NamedTemporaryFile(delete=False) as temp_file: | |
| shutil.copyfileobj(image_file, temp_file) | |
| image = Image.open(temp_file.name) | |
| image = np.array(image) | |
| logging.info('Starting process_images_and_statements') | |
| # Generate the image caption | |
| generated_caption = caption_image(image) | |
| # Match the statements | |
| matched_statements = match_statements(image, statements) | |
| os.unlink(temp_file.name) # Remove the temporary file | |
| return generated_caption, matched_statements | |
| # Define Gradio interface | |
| image_input = gr.inputs.Image(type="numpy", label="Upload Image") | |
| outputs = [ | |
| gr.outputs.Image(type="pil", label="Annotated Image"), | |
| gr.outputs.Textbox(label="Matched Statements"), | |
| ] | |
| iface = gr.Interface( | |
| fn=process_images_and_statements, | |
| inputs=image_input, | |
| outputs=outputs, | |
| title="Image Captioning and Matching", | |
| description="Upload an image to generate a caption for the image and match the statements.", | |
| theme='sudeepshouche/minimalist' | |
| ) | |
| # Launch Gradio app | |
| iface.launch(debug=True) | |
| # Launch Gradio app | |
| iface.launch(debug=True) | |