Update app.py

app.py

@@ -1,29 +1,71 @@
-from PIL import Image
-import requests
-from io import BytesIO
-from transformers import pipeline
-import streamlit as st
-def predict(image):
-    type_food = oracle(image, "What type of food is this?")
-    cal_est = oracle(image, "About how many calories are in this meal?")
-    guess1, guess2 = cal_est[0]['answer'], cal_est[1]['answer']
-    return f"This is {type_food[0]['answer']}. I estimate this to contain {min(guess1, guess2)}-{max(guess1, guess2)} calories"
-oracle = pipeline(model="dandelin/vilt-b32-finetuned-vqa")
-
-def main():
-    st.title("Image Question Answering App")
-    st.write("Upload an image and ask a question to get answers!")
+import torch 
+import re 
+import gradio as gr
+from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel 
+
+device='cpu'
+encoder_checkpoint = "nlpconnect/vit-gpt2-image-captioning"
+decoder_checkpoint = "nlpconnect/vit-gpt2-image-captioning"
+model_checkpoint = "nlpconnect/vit-gpt2-image-captioning"
+feature_extractor = ViTFeatureExtractor.from_pretrained(encoder_checkpoint)
+tokenizer = AutoTokenizer.from_pretrained(decoder_checkpoint)
+model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint).to(device)
+
+
+def predict(image,max_length=64, num_beams=3):
+  image = image.convert('RGB')
+  image = feature_extractor(image, return_tensors="pt").pixel_values.to(device)
+  clean_text = lambda x: x.replace('<|endoftext|>','').split('\n')[0]
+  caption_ids = model.generate(image, max_length = max_length)[0]
+  caption_text = clean_text(tokenizer.decode(caption_ids))
+  return caption_text 
+
+
+
+input = gr.inputs.Image(label="Upload any Image", type = 'pil', optional=True)
+output = gr.outputs.Textbox(type="auto",label="Captions")
+examples = [f"example{i}.jpg" for i in range(1,7)]
+
+title = "Image Captioning "
+description = "Made by : shreyasdixit.tech"
+interface = gr.Interface(
+            
+        fn=predict,
+        description=description,
+        inputs = input,
+        theme="grass",
+        outputs=output,
+        examples = examples,
+        title=title,
+    )
+interface.launch(debug=True)
+
+# from PIL import Image
+# import requests
+# from io import BytesIO
+# from transformers import pipeline
+# import streamlit as st
+# def predict(image):
+#     type_food = oracle(image, "What type of food is this?")
+#     cal_est = oracle(image, "About how many calories are in this meal?")
+#     guess1, guess2 = cal_est[0]['answer'], cal_est[1]['answer']
+#     return f"This is {type_food[0]['answer']}. I estimate this to contain {min(guess1, guess2)}-{max(guess1, guess2)} calories"
+# oracle = pipeline(model="dandelin/vilt-b32-finetuned-vqa")
+
+# def main():
+#     st.title("Image Question Answering App")
+#     st.write("Upload an image and ask a question to get answers!")
     
-    oracle = pipeline(model="dandelin/vilt-b32-finetuned-vqa")
+#     oracle = pipeline(model="dandelin/vilt-b32-finetuned-vqa")
 
-    # File uploader for image
-    uploaded_image = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
+#     # File uploader for image
+#     uploaded_image = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
     
-    if uploaded_image is not None:
-        image = Image.open(uploaded_image)
-        st.image(image, caption="Uploaded Image", use_column_width=True)
-        response = predict(image)
-        st.write(response)
+#     if uploaded_image is not None:
+#         image = Image.open(uploaded_image)
+#         st.image(image, caption="Uploaded Image", use_column_width=True)
+#         response = predict(image)
+#         st.write(response)
             
-if __name__ == "__main__":
-    main()
+# if __name__ == "__main__":
+#     main()