Bring up to date with working github copy

https://github.com/brianjking/image-caption-textmatch/commit/4ec79f5cccf8269f713977583c4bda88a34c8ca3

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
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor
 
 # 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)
-tokenizer_caption = AutoTokenizer.from_pretrained("microsoft/git-large-r-textcaps")
-model_caption = AutoModelForCausalLM.from_pretrained("microsoft/git-large-r-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")
 
 # List of statements for Image-Text Matching
 statements = [
@@ -25,41 +25,51 @@ statements = [
     'promotes alcohol use as a "rite of passage" to adulthood',
 ]
 
-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):
+# Function to compute ITM scores for the combined text input (caption + statement)
+def compute_itm_score(image, combined_text):
     pil_image = Image.fromarray(image.astype('uint8'), 'RGB')
     img = vis_processors["eval"](pil_image.convert("RGB")).unsqueeze(0).to(device)
-    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
+    # 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
+
+
+
+
 
-# 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
+
+
+
+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
 
 # Main function to perform image captioning and image-text matching
 def process_images_and_statements(image):
-    # Generate image captions using TextCaps
-    captions = generate_image_captions()
+    # Generate image caption for the uploaded image using git-large-r-textcaps
+    caption = generate_caption(git_processor_large_textcaps, git_model_large_textcaps, image)
+
+    # Initialize an empty list to store the results
+    results = []
     
-    # Compute ITM scores for predefined statements using LAVIS
-    itm_scores = compute_itm_scores(image)
+    # Loop through each predefined statement
+    for statement in statements:
+        # Concatenate the caption with the statement
+        combined_text = caption + " " + statement
+
+        # 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 image captions and ITM scores into the output
-    output = "Image Captions:\n" + captions + "\n\nITM Scores:\n" + itm_scores
+    # Combine the results and return them
+    output = "\n".join(results)
     return output
 
 # Gradio interface
@@ -67,4 +77,4 @@ image_input = gr.inputs.Image()
 output = gr.outputs.Textbox(label="Results")
 
 iface = gr.Interface(fn=process_images_and_statements, inputs=image_input, outputs=output, title="Image Captioning and Image-Text Matching")
-iface.launch()
+iface.launch()