Update pages/Captionize.py

pages/Captionize.py

@@ -1,8 +1,30 @@
 import torch 
 import re 
-import gradio as gr
+import streamlit as st
 from transformers import AutoTokenizer, ViTFeatureExtractor, VisionEncoderDecoderModel 
 
+
+st.set_page_config(page_title="Captionize")
+
+st.title("🤖 Captionize")
+st.subheader("Generate Captions for your Image...")
+
+st.sidebar.image('./csv_analysis.png',width=300, use_column_width=True)
+
+# Applying Styling
+st.markdown("""
+<style>
+div.stButton > button:first-child {
+    background-color: #0099ff;
+    color:#ffffff;
+}
+div.stButton > button:hover {
+    background-color: #00ff00;
+    color:#FFFFFF;
+    }
+</style>""", unsafe_allow_html=True)
+
+
 device='cpu'
 encoder_checkpoint = "nlpconnect/vit-gpt2-image-captioning"
 decoder_checkpoint = "nlpconnect/vit-gpt2-image-captioning"
@@ -18,24 +40,14 @@ def predict(image,max_length=64, num_beams=4):
   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)
+  return caption_text
+
+pic = st.file_uploader(label="Please upload any Image here 😎",type=['png', 'jpeg', 'jpg'], help="Only 'png', 'jpeg' or 'jpg' formats allowed")
+
+
+button = st.button("Generate Caption")
+
+if button:
+    # Get Response
+    caption =  predict(pic)
+    st.write(caption)