updated code

app.py

@@ -1,59 +1,66 @@
 import streamlit as st
 import io
 
-# x = st.slider("Select a value")
-# st.write(x, "squared is", x * x)
-
-st.title("Image Captioning Demo from Redcaps")
-st.sidebar.markdown(
-    """
-    Image Captioning Model from VirTex trained on Redcaps
-    """
-)
-
-with st.spinner("Loading Model"):
-    from model import *
-    sample_images = glob.glob("./samples/*.jpg")
-    download_files()
-    virtexModel = VirTexModel()
-    imageLoader = ImageLoader()
-
-random_image = get_rand_img(sample_images)
-
-st.sidebar.title("Select a sample image")
-sample_image = st.sidebar.selectbox(
-    "",
-    sample_images
-)
-
-if st.sidebar.button("Random Sample Image"):
-    random_image = get_rand_img(sample_images)
-    sample_image = None
+# st.title("Image Captioning Demo from Redcaps")
+# st.sidebar.markdown(
+#     """
+#     Image Captioning Model from VirTex trained on Redcaps
+#     """
+# )
+
+# with st.spinner("Loading Model"):
+#     from model import *
+#     sample_images = glob.glob("./samples/*.jpg")
+#     download_files()
+#     virtexModel = VirTexModel()
+#     imageLoader = ImageLoader()
+
+# random_image = get_rand_img(sample_images)
+
+# st.sidebar.title("Select a sample image")
+# sample_image = st.sidebar.selectbox(
+#     "",
+#     sample_images
+# )
+
+# if st.sidebar.button("Random Sample Image"):
+#     random_image = get_rand_img(sample_images)
+#     sample_image = None
     
-uploaded_image = None
-with st.sidebar.form("file-uploader-form", clear_on_submit=True):
-    uploaded_file = st.file_uploader("Choose a file")
-    submitted = st.form_submit_button("Submit")
-    if uploaded_file is not None and submitted:
-        uploaded_image = Image.open(io.BytesIO(uploaded_file.get_values()))
-
-if uploaded_image is None and submitted:
-    st.write("Please select a file to upload")
-
-else:
-    image_file = sample_image if sample_image is not None else random_image
+# uploaded_image = None
+# with st.sidebar.form("file-uploader-form", clear_on_submit=True):
+#     uploaded_file = st.file_uploader("Choose a file")
+#     submitted = st.form_submit_button("Submit")
+#     if uploaded_file is not None and submitted:
+#         uploaded_image = Image.open(io.BytesIO(uploaded_file.get_values()))
+
+# if uploaded_image is None and submitted:
+#     st.write("Please select a file to upload")
+
+# else:
+#     image_file = sample_image if sample_image is not None else random_image
     
-    image = uploaded_image if uploaded_image is not None else Image.open()
+#     image = uploaded_image if uploaded_image is not None else Image.open()
     
-    image_dict = imageLoader.transform(image)
+#     image_dict = imageLoader.transform(image)
     
-    show.image(st.image(image_dict["image"]), "Target Image")
+#     show.image(st.image(image_dict["image"]), "Target Image")
     
-    with st.spinner("Generating Caption"):
-        subreddit, caption = virtexModel.predict(image_dict)
-        st.header("Predicted Caption:\n\n")
-        st.subheader(f"Subreddit: {subreddit}\n")
-        st.subheader(f"Caption: {caption}\n")
+#     with st.spinner("Generating Caption"):
+#         subreddit, caption = virtexModel.predict(image_dict)
+#         st.header("Predicted Caption:\n\n")
+#         st.subheader(f"Subreddit: {subreddit}\n")
+#         st.subheader(f"Caption: {caption}\n")
         
-    image.close()
+#     image.close()
+
+from model import *
+download_files()
+sample_images = get_samples()
+v, il = VirTexModel(), ImageLoader()
 
+for s in sample_images:
+    subreddit, caption = v.predict(il.load(s))
+    print("=====================")
+    print(subreddit)
+    print(caption)

model.py

@@ -24,12 +24,12 @@ class ImageLoader():
         self.transformer = torchvision.transforms.Compose([torchvision.transforms.Resize(256),
                                torchvision.transforms.CenterCrop(224),
                                torchvision.transforms.ToTensor()])
-    def load(self, im_path, prompt = ""):
+    def load(self, im_path):
         im = torch.FloatTensor(self.transformer(Image.open(im_path))).unsqueeze(0)
-        return {"image": im, "decode_prompt": prompt}
-    def transform(self, image, prompt = ""):
+        return {"image": im}
+    def transform(self, image):
         im = torch.FloatTensor(self.transformer(image)).unsqueeze(0)
-        return {"image": im, "decode_prompt": prompt}
+        return {"image": im}
 
 class VirTexModel():
     def __init__(self):
@@ -51,6 +51,7 @@ class VirTexModel():
         
         is_valid_subreddit = False
         subreddit, rest_of_caption = "", ""
+        image_dict["decode_prompt"] = subreddit_tokens
         while not is_valid_subreddit:
             
             with torch.no_grad():