Update app.py

app.py

@@ -1,7 +1,6 @@
 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
@@ -12,7 +11,7 @@ model_itm, vis_processors, text_processors = load_model_and_preprocess("blip2_im
 
 # 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")
+model_caption = AutoModelForCausalLM.from_pretrained("microsoft/git-large-r-textcaps").to(device)
 
 # List of statements for Image-Text Matching
 statements = [
@@ -25,41 +24,38 @@ 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):
-    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
-
 # 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)
+def generate_image_captions(image):
+    inputs = tokenizer_caption(image, return_tensors="pt", padding=True, truncation=True).to(device)
     outputs = model_caption.generate(**inputs)
     caption = tokenizer_caption.decode(outputs[0], skip_special_tokens=True)
-    return prompt + " " + caption
+    return caption
 
 # Main function to perform image captioning and image-text matching
 def process_images_and_statements(image):
+    pil_image = Image.fromarray(image.astype('uint8'), 'RGB')
+    img = vis_processors["eval"](pil_image.convert("RGB")).unsqueeze(0).to(device)
+    
     # Generate image captions using TextCaps
-    captions = generate_image_captions()
+    captions = generate_image_captions(pil_image)
     
-    # Compute ITM scores for predefined statements using LAVIS
-    itm_scores = compute_itm_scores(image)
+    # Convert caption to the format expected by the ITM model
+    txt = text_processors["eval"](captions)
     
-    # Combine image captions and ITM scores into the output
-    output = "Image Captions:\n" + captions + "\n\nITM Scores:\n" + itm_scores
+    # Compute ITM scores for predefined statements
+    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()
+    
+    results = [f'Image Caption: "{captions}" with a matching probability of {score:.3%}']
+    for statement in statements:
+        txt = text_processors["eval"](statement)
+        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 combination of image, caption ("{captions}"), and statement ("{statement}") is matched with a probability of {score:.3%}'
+        results.append(result_text)
+    output = "\n".join(results)
     return output
 
 # Gradio interface
@@ -67,4 +63,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()