Update app.py

app.py

@@ -1,71 +1,30 @@
-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!")
+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()