Update app.py

app.py

@@ -43,17 +43,17 @@ def compute_itm_scores(image):
     return output
 
 # Function to generate image captions using TextCaps
-def generate_image_captions(image):
-    pil_image = Image.fromarray(image.astype('uint8'), 'RGB')
-    inputs = tokenizer_caption(pil_image, return_tensors="pt", padding=True, truncation=True)
+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])
-    return caption
+    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 captions using TextCaps
-    captions = generate_image_captions(image)
+    captions = generate_image_captions()
     
     # Compute ITM scores for predefined statements using LAVIS
     itm_scores = compute_itm_scores(image)
@@ -64,7 +64,4 @@ def process_images_and_statements(image):
 
 # Gradio interface
 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()
+output = gr.outputs.Textbox(label