Update app.py

app.py

@@ -4,15 +4,15 @@ 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
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # 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")
+tokenizer_caption = AutoTokenizer.from_pretrained("microsoft/git-large-r-textcaps")
+model_caption = AutoModelForCausalLM.from_pretrained("microsoft/git-large-r-textcaps")
 
 # List of statements for Image-Text Matching
 statements = [
@@ -25,51 +25,41 @@ statements = [
     'promotes alcohol use as a "rite of passage" to adulthood',
 ]
 
-# Function to compute ITM scores for the combined text input (caption + statement)
-def compute_itm_score(image, combined_text):
+txts = [text_processors["eval"](statement) for statement in statements]
+
+# Function to compute Image-Text Matching (ITM) scores for all statements
+def compute_itm_scores(image):
     pil_image = Image.fromarray(image.astype('uint8'), 'RGB')
     img = vis_processors["eval"](pil_image.convert("RGB")).unsqueeze(0).to(device)
-    # Pass the combined_text string directly to model_itm
-    itm_output = model_itm({"image": img, "text_input": combined_text}, match_head="itm")
-    itm_scores = torch.nn.functional.softmax(itm_output, dim=1)
-    score = itm_scores[:, 1].item()
-    return score
-
-
-
-
-
-
-
+    results = []
+    for i, statement in enumerate(statements):
+        txt = txts[i]
+        itm_output = model_itm({"image": img, "text_input": txt}, match_head="itm")
+        itm_scores = torch.nn.functional.softmax(itm_output, dim=1)
+        score = itm_scores[:, 1].item()
+        result_text = f'The image and "{statement}" are matched with a probability of {score:.3%}'
+        results.append(result_text)
+    output = "\n".join(results)
+    return output
 
-def generate_caption(processor, model, image):
-    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]
-    return generated_caption
+# Function to generate image captions using TextCaps
+def generate_image_captions():
+    prompt = "A photo of"
+    inputs = tokenizer_caption(prompt, return_tensors="pt", padding=True, truncation=True)
+    outputs = model_caption.generate(**inputs)
+    caption = tokenizer_caption.decode(outputs[0], skip_special_tokens=True)
+    return prompt + " " + caption
 
 # Main function to perform image captioning and image-text matching
 def process_images_and_statements(image):
-    # Generate image caption for the uploaded image using git-large-r-textcaps
-    caption = generate_caption(git_processor_large_textcaps, git_model_large_textcaps, image)
+    # Generate image captions using TextCaps
+    captions = generate_image_captions()
 
-    # Initialize an empty list to store the results
-    results = []
-    
-    # Loop through each predefined statement
-    for statement in statements:
-        # Concatenate the caption with the statement
-        combined_text = caption + " " + statement
+    # Compute ITM scores for predefined statements using LAVIS
+    itm_scores = compute_itm_scores(image)
 
-        # Compute ITM score for the combined text and the image
-        itm_score = compute_itm_score(image, combined_text)
-        
-        # Store the result
-        result_text = f'The image and "{combined_text}" are matched with a probability of {itm_score:.3%}'
-        results.append(result_text)
-    
-    # Combine the results and return them
-    output = "\n".join(results)
+    # Combine image captions and ITM scores into the output
+    output = "Image Captions:\n" + captions + "\n\nITM Scores:\n" + itm_scores
     return output
 
 # Gradio interface