Create app.py

GitHub Repository: https://github.com/CapstoneProjectimagecaptioning/image_captioning_transformer

app.py

@@ -0,0 +1,54 @@
+import torch
+import gradio as gr
+from transformers import AutoTokenizer, ViTImageProcessor, VisionEncoderDecoderModel
+
+device = 'cpu'
+
+# Load the pretrained model, feature extractor, and tokenizer
+model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning").to(device)
+feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
+
+def predict(image, max_length=64, num_beams=4):
+    # Process the input image
+    image = image.convert('RGB')
+    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device)
+    
+    # Generate the caption
+    caption_ids = model.generate(pixel_values, max_length=max_length, num_beams=num_beams)[0]
+    
+    # Decode and clean the generated caption
+    caption = tokenizer.decode(caption_ids, skip_special_tokens=True)
+    return caption
+
+css = '''
+h1#title {
+  text-align: center;
+}
+h3#header {
+  text-align: center;
+}
+img#overview {
+  max-width: 800px;
+  max-height: 600px;
+}
+img#style-image {
+  max-width: 1000px;
+  max-height: 600px;
+}
+'''
+
+demo = gr.Blocks(css=css)
+
+with demo:
+    gr.Markdown('''<h1 id="title">Automated Image Captioning Using Generative AI: A Transformer based approach 🖼️</h1>''')
+    gr.Markdown('Contributed by : Premanth Alahari, Charan Gudivada')
+    
+    with gr.Column():
+        input_image = gr.Image(label="Upload your Image", type='pil')
+        output_caption = gr.Textbox(label="Generated Caption")
+        
+    btn = gr.Button("Generate Caption")
+    btn.click(fn=predict, inputs=input_image, outputs=output_caption)
+
+demo.launch()