Update app.py

app.py

@@ -1,38 +1,41 @@
 from PIL import Image
-from transformers import VisionEncoderDecoderModel , ViTFeatureExtractor , PreTrainedTokenizerFast
+from transformers import VisionEncoderDecoderModel, ViTFeatureExtractor, PreTrainedTokenizerFast
 import gradio as gr
 
-model = VisionEncoderDecoderModel.from_pretrained("ydshieh/vit-gpt2-coco-en")
-vit_feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch32-224-in21k")
-tokenizer = PreTrainedTokenizerFast.from_pretrained("distilgpt2")
-
+# Load the model and processor
+model = VisionEncoderDecoderModel.from_pretrained("microsoft/git-base")
+feature_extractor = ViTFeatureExtractor.from_pretrained("microsoft/git-base")
+tokenizer = PreTrainedTokenizerFast.from_pretrained("microsoft/git-base")
 
+# Define the captioning function
 def caption_images(image):
-  pixel_values = vit_feature_extractor(images=image,return_tensors="pt").pixel_values
-  encoder_outputs = model.generate(pixel_values.to('cpu'),num_beams=5)
-  generated_sentence = tokenizer.batch_decode(encoder_outputs,skip_special_tokens=True)
-
-  return (generated_sentence[0].strip())
-
-
+    # Preprocess the image
+    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
+    # Generate captions
+    encoder_outputs = model.generate(pixel_values.to('cpu'), num_beams=5)
+    generated_sentence = tokenizer.batch_decode(encoder_outputs, skip_special_tokens=True)
+    return generated_sentence[0].strip()
+
+# Define Gradio interface components
 inputs = [
-    gr.components.Image(type='pil',label='Original Image')
+    gr.inputs.Image(type='pil', label='Original Image')
 ]
 
 outputs = [
-    gr.components.Textbox(label='Caption')
+    gr.outputs.Textbox(label='Caption')
 ]
 
-title = "Simple Image captioning Application"
+# Define Gradio app properties
+title = "Simple Image Captioning Application"
 description = "Upload an image to see the caption generated"
-example =['messi.jpg']
+example = ['messi.jpg']  # Replace with a valid path to an example image
 
+# Create and launch the Gradio interface
 gr.Interface(
-    caption_images,
-    inputs,
-    outputs,
+    fn=caption_images,
+    inputs=inputs,
+    outputs=outputs,
     title=title,
-    description = description,
-    examples = example,
+    description=description,
+    examples=example,
 ).launch(debug=True)
-