prompt ranking working

streamlit_app.py → app.py

@@ -4,9 +4,6 @@ import requests
 import streamlit as st
 from clip_model import ClipModel
 
-from session_state import SessionState, get_state
-from images_mocker import ImagesMocker
-
 from PIL import Image
 
 IMAGES_LINKS = ["https://cdn.pixabay.com/photo/2014/10/13/21/34/clipper-487503_960_720.jpg",
@@ -32,24 +29,34 @@ def load_image_from_url(url: str) -> Image.Image:
 def load_model() -> ClipModel:
     return ClipModel()
 
+def init_state():
+    if "images" not in st.session_state:
+        st.session_state.images = None
+    if "prompts" not in st.session_state:
+        st.session_state.prompts = None
+    if "predictions" not in st.session_state:
+        st.session_state.predictions = None
+    if "default_text_input" not in st.session_state:
+        st.session_state.default_text_input = None
+
 
-def limit_number_images(state: SessionState):
+def limit_number_images():
     """When moving between tasks sometimes the state of images can have too many samples"""
-    if state.images is not None and len(state.images) > 1:
-        state.images = [state.images[0]]
+    if st.session_state.images is not None and len(st.session_state.images) > 1:
+        st.session_state.images = [st.session_state.images[0]]
 
 
-def limit_number_prompts(state: SessionState):
+def limit_number_prompts():
     """When moving between tasks sometimes the state of prompts can have too many samples"""
-    if state.prompts is not None and len(state.prompts) > 1:
-        state.prompts = [state.prompts[0]]
+    if st.session_state.prompts is not None and len(st.session_state.prompts) > 1:
+        st.session_state.prompts = [st.session_state.prompts[0]]
 
 
-def is_valid_prediction_state(state: SessionState) -> bool:
-    if state.images is None or len(state.images) < 1:
+def is_valid_prediction_state() -> bool:
+    if st.session_state.images is None or len(st.session_state.images) < 1:
         st.error("Choose at least one image before predicting")
         return False
-    if state.prompts is None or len(state.prompts) < 1:
+    if st.session_state.prompts is None or len(st.session_state.prompts) < 1:
         st.error("Write at least one prompt before predicting")
         return False
     return True
@@ -97,16 +104,16 @@ class Sections:
         st.markdown("### Try OpenAI's CLIP model in your browser")
         st.markdown(" ")
         st.markdown(" ")
-        with st.beta_expander("What is CLIP?"):
+        with st.expander("What is CLIP?"):
             st.markdown("CLIP is a machine learning model that computes similarity between text "
                         "(also called prompts) and images. It has been trained on a dataset with millions of diverse"
                         " image-prompt pairs, which allows it to generalize to unseen examples."
                         " <br /> Check out [OpenAI's blogpost](https://openai.com/blog/clip/) for more details",
                         unsafe_allow_html=True)
-            col1, col2 = st.beta_columns(2)
+            col1, col2 = st.columns(2)
             col1.image("https://openaiassets.blob.core.windows.net/$web/clip/draft/20210104b/overview-a.svg")
             col2.image("https://openaiassets.blob.core.windows.net/$web/clip/draft/20210104b/overview-b.svg")
-        with st.beta_expander("What can CLIP do?"):
+        with st.expander("What can CLIP do?"):
             st.markdown("#### Prompt ranking")
             st.markdown("Given different prompts and an image CLIP will rank the different prompts based on how well they describe the image")
             st.markdown("#### Image ranking")
@@ -118,7 +125,7 @@ class Sections:
         st.markdown(" ")
 
     @staticmethod
-    def image_uploader(state: SessionState, accept_multiple_files: bool):
+    def image_uploader(accept_multiple_files: bool):
         uploaded_images = st.file_uploader("Upload image", type=[".png", ".jpg", ".jpeg"],
                                            accept_multiple_files=accept_multiple_files)
         if (not accept_multiple_files and uploaded_images is not None) or (accept_multiple_files and len(uploaded_images) >= 1):
@@ -129,125 +136,123 @@ class Sections:
                 pil_image = Image.open(uploaded_image)
                 pil_image = preprocess_image(pil_image)
                 images.append(pil_image)
-            state.images = images
+            st.session_state.images = images
 
 
     @staticmethod
-    def image_picker(state: SessionState, default_text_input: str):
-        col1, col2, col3 = st.beta_columns(3)
+    def image_picker(default_text_input: str):
+        col1, col2, col3 = st.columns(3)
         with col1:
             default_image_1 = load_image_from_url("https://cdn.pixabay.com/photo/2014/10/13/21/34/clipper-487503_960_720.jpg")
             st.image(default_image_1, use_column_width=True)
             if st.button("Select image 1"):
-                state.images = [default_image_1]
-                state.default_text_input = default_text_input
+                st.session_state.images = [default_image_1]
+                st.session_state.default_text_input = default_text_input
         with col2:
             default_image_2 = load_image_from_url("https://cdn.pixabay.com/photo/2019/11/11/14/30/zebra-4618513_960_720.jpg")
             st.image(default_image_2, use_column_width=True)
             if st.button("Select image 2"):
-                state.images = [default_image_2]
-                state.default_text_input = default_text_input
+                st.session_state.images = [default_image_2]
+                st.session_state.default_text_input = default_text_input
         with col3:
             default_image_3 = load_image_from_url("https://cdn.pixabay.com/photo/2016/11/15/16/24/banana-1826760_960_720.jpg")
             st.image(default_image_3, use_column_width=True)
             if st.button("Select image 3"):
-                state.images = [default_image_3]
-                state.default_text_input = default_text_input
+                st.session_state.images = [default_image_3]
+                st.session_state.default_text_input = default_text_input
 
     @staticmethod
-    def dataset_picker(state: SessionState):
-        columns = st.beta_columns(5)
-        state.dataset = load_default_dataset()
+    def dataset_picker():
+        columns = st.columns(5)
+        st.session_state.dataset = load_default_dataset()
         image_idx = 0
         for col in columns:
-            col.image(state.dataset[image_idx])
+            col.image(st.session_state.dataset[image_idx])
             image_idx += 1
-            col.image(state.dataset[image_idx])
+            col.image(st.session_state.dataset[image_idx])
             image_idx += 1
         if st.button("Select random dataset"):
-            state.images = state.dataset
-            state.default_text_input = "A sign that says 'SLOW DOWN'"
+            st.session_state.images = st.session_state.dataset
+            st.session_state.default_text_input = "A sign that says 'SLOW DOWN'"
 
     @staticmethod
-    def prompts_input(state: SessionState, input_label: str, prompt_prefix: str = ''):
+    def prompts_input(input_label: str, prompt_prefix: str = ''):
         raw_text_input = st.text_input(input_label,
-                                    value=state.default_text_input if state.default_text_input is not None else "")
-        state.is_default_text_input = raw_text_input == state.default_text_input
+                                    value=st.session_state.default_text_input if st.session_state.default_text_input is not None else "")
+        st.session_state.is_default_text_input = raw_text_input == st.session_state.default_text_input
         if raw_text_input:
-            state.prompts = [prompt_prefix + class_name for class_name in raw_text_input.split(";") if len(class_name) > 1]
+            st.session_state.prompts = [prompt_prefix + class_name for class_name in raw_text_input.split(";") if len(class_name) > 1]
 
     @staticmethod
-    def single_image_input_preview(state: SessionState):
+    def single_image_input_preview():
         st.markdown("### Preview")
-        col1, col2 = st.beta_columns([1, 2])
+        col1, col2 = st.columns([1, 2])
         with col1:
             st.markdown("Image to classify")
-            if state.images is not None:
-                st.image(state.images[0], use_column_width=True)
+            if st.session_state.images is not None:
+                st.image(st.session_state.images[0], use_column_width=True)
             else:
                 st.warning("Select an image")
 
         with col2:
             st.markdown("Labels to choose from")
-            if state.prompts is not None:
-                for prompt in state.prompts:
+            if st.session_state.prompts is not None:
+                for prompt in st.session_state.prompts:
                     st.markdown(f"* {prompt}")
-                if len(state.prompts) < 2:
+                if len(st.session_state.prompts) < 2:
                     st.warning("At least two prompts/classes are needed")
             else:
                 st.warning("Enter the prompts/classes to classify from")
 
     @staticmethod
-    def multiple_images_input_preview(state: SessionState):
+    def multiple_images_input_preview():
         st.markdown("### Preview")
         st.markdown("Images to classify")
-        col1, col2, col3 = st.beta_columns(3)
-        if state.images is not None:
-            for idx, image in enumerate(state.images):
-                if idx < len(state.images) / 2:
-                    col1.image(state.images[idx], use_column_width=True)
+        col1, col2, col3 = st.columns(3)
+        if st.session_state.images is not None:
+            for idx, image in enumerate(st.session_state.images):
+                if idx < len(st.session_state.images) / 2:
+                    col1.image(st.session_state.images[idx], use_column_width=True)
                 else:
-                    col2.image(state.images[idx], use_column_width=True)
-            if len(state.images) < 2:
+                    col2.image(st.session_state.images[idx], use_column_width=True)
+            if len(st.session_state.images) < 2:
                 col2.warning("At least 2 images required")
         else:
             col1.warning("Select an image")
 
         with col3:
             st.markdown("Query prompt")
-            if state.prompts is not None:
-                for prompt in state.prompts:
+            if st.session_state.prompts is not None:
+                for prompt in st.session_state.prompts:
                     st.write(prompt)
             else:
                 st.warning("Enter the prompt to classify")
 
     @staticmethod
-    def classification_output(state: SessionState, model: ClipModel):
+    def classification_output(model: ClipModel):
         # Possible way of customize this https://discuss.streamlit.io/t/st-button-in-a-custom-layout/2187/2
-        if st.button("Predict") and is_valid_prediction_state(state):  # PREDICT 🚀
+        if st.button("Predict") and is_valid_prediction_state():  # PREDICT 🚀
             with st.spinner("Predicting..."):
                 
                 st.markdown("### Results")
-                # st.write(clip_response)
-                if len(state.images) == 1:
-                    scores = model.compute_prompts_probabilities(state.images[0], state.prompts)
-                    scored_prompts = [(prompt, score) for prompt, score in zip(state.prompts, scores)]
-                    st.json(scores)
+                if len(st.session_state.images) == 1:
+                    scores = model.compute_prompts_probabilities(st.session_state.images[0], st.session_state.prompts)
+                    scored_prompts = [(prompt, score) for prompt, score in zip(st.session_state.prompts, scores)]
                     sorted_scored_prompts = sorted(scored_prompts, key=lambda x: x[1], reverse=True)
                     for prompt, probability in sorted_scored_prompts:
                         percentage_prob = int(probability * 100)
                         st.markdown(
-                            f"### ![prob](https://progress-bar.dev/{percentage_prob}/?width=200) &nbsp &nbsp {prompt}")
-                elif len(state.prompts) == 1:
-                    st.markdown(f"### {state.prompts[0]}")
+                            f"### ![prob](https://progress-bar.dev/{percentage_prob}/?width=200) {prompt}")
+                elif len(st.session_state.prompts) == 1:
+                    st.markdown(f"### {st.session_state.prompts[0]}")
                     
-                    scores = model.compute_prompts_probabilities(state.images[0], state.prompts)
-                    scored_images = [(image, score) for image, score in zip(state.images, scores)]
+                    scores = model.compute_images_probabilities(st.session_state.images, st.session_state.prompts[0])
                     st.json(scores)
-                    sorted_scored_images = sorted(scored_prompts, key=lambda x: x[1], reverse=True)
+                    scored_images = [(image, score) for image, score in zip(st.session_state.images, scores)]
+                    sorted_scored_images = sorted(scored_images, key=lambda x: x[1], reverse=True)
 
                     for image, probability in sorted_scored_images[:5]:
-                        col1, col2 = st.beta_columns([1, 3])
+                        col1, col2 = st.columns([1, 3])
                         col1.image(image, use_column_width=True)
                         percentage_prob = int(probability * 100)
                         col2.markdown(f"### ![prob](https://progress-bar.dev/{percentage_prob}/?width=200)")
@@ -269,47 +274,45 @@ class Sections:
 
 
 Sections.header()
-col1, col2 = st.beta_columns([1, 2])
+col1, col2 = st.columns([1, 2])
 col1.markdown(" "); col1.markdown(" ")
 col1.markdown("#### Task selection")
 task_name: str = col2.selectbox("", options=["Prompt ranking", "Image ranking", "Image classification"])
 st.markdown("<br>", unsafe_allow_html=True)
-
+init_state()
 model = load_model()
-session_state = get_state()
 if task_name == "Image classification":
-    Sections.image_uploader(session_state, accept_multiple_files=False)
-    if session_state.images is None:
+    Sections.image_uploader(accept_multiple_files=False)
+    if st.session_state.images is None:
         st.markdown("or choose one from")
-        Sections.image_picker(session_state, default_text_input="banana; boat; bird")
+        Sections.image_picker(default_text_input="banana; boat; bird")
     input_label = "Enter the classes to chose from separated by a semi-colon. (f.x. `banana; boat; honesty; apple`)"
-    Sections.prompts_input(session_state, input_label, prompt_prefix='A picture of a ')
-    limit_number_images(session_state)
-    Sections.single_image_input_preview(session_state)
-    Sections.classification_output(session_state, model)
+    Sections.prompts_input(input_label, prompt_prefix='A picture of a ')
+    limit_number_images()
+    Sections.single_image_input_preview()
+    Sections.classification_output(model)
 elif task_name == "Prompt ranking":
-    Sections.image_uploader(session_state, accept_multiple_files=False)
-    if session_state.images is None:
+    Sections.image_uploader(accept_multiple_files=False)
+    if st.session_state.images is None:
         st.markdown("or choose one from")
-        Sections.image_picker(session_state, default_text_input="A calm afternoon in the Mediterranean; "
+        Sections.image_picker(default_text_input="A calm afternoon in the Mediterranean; "
                                                                 "A beautiful creature;"
                                                                 " Something that grows in tropical regions")
     input_label = "Enter the prompts to choose from separated by a semi-colon. " \
                   "(f.x. `An image that inspires; A feeling of loneliness; joyful and young; apple`)"
-    Sections.prompts_input(session_state, input_label)
-    limit_number_images(session_state)
-    Sections.single_image_input_preview(session_state)
-    Sections.classification_output(session_state, model)
+    Sections.prompts_input(input_label)
+    limit_number_images()
+    Sections.single_image_input_preview()
+    Sections.classification_output(model)
 elif task_name == "Image ranking":
-    Sections.image_uploader(session_state, accept_multiple_files=True)
-    if session_state.images is None or len(session_state.images) < 2:
+    Sections.image_uploader(accept_multiple_files=True)
+    if st.session_state.images is None or len(st.session_state.images) < 2:
         st.markdown("or use this random dataset")
-        Sections.dataset_picker(session_state)
-    Sections.prompts_input(session_state, "Enter the prompt to query the images by")
-    limit_number_prompts(session_state)
-    Sections.multiple_images_input_preview(session_state)
-    Sections.classification_output(session_state, model)
+        Sections.dataset_picker()
+    Sections.prompts_input("Enter the prompt to query the images by")
+    limit_number_prompts()
+    Sections.multiple_images_input_preview()
+    Sections.classification_output(model)
 
 st.markdown("<br><br><br><br>Made by [@JavierFnts](https://twitter.com/JavierFnts) | [How was CLIP Playground built?](https://twitter.com/JavierFnts/status/1363522529072214019)"
             "", unsafe_allow_html=True)
-session_state.sync()

clip_model.py

@@ -9,7 +9,7 @@ class ClipModel:
         ['RN50', 'RN101', 'RN50x4', 'RN50x16', 'RN50x64', 'ViT-B/32',
          'ViT-B/16', 'ViT-L/14', 'ViT-L/14@336px']
         """
-        self.model, self.img_preprocess = clip.load(model_name)
+        self._model, self._img_preprocess = clip.load(model_name)
     
     def predict(self, images: list[Image], prompts: list[str]) -> dict:
         if len(images) == 1:
@@ -19,43 +19,40 @@ class ClipModel:
         else:
             raise ValueError('Either images or prompts must be a single element')
     
-    def compute_prompts_probabilities(self, image: Image, prompts: list[str]) -> dict[str, float]:
-        preprocessed_image = self.img_preprocess(image).unsqueeze(0)
+    def compute_prompts_probabilities(self, image: Image, prompts: list[str]) -> list[float]:
+        preprocessed_image = self._img_preprocess(image).unsqueeze(0)
         tokenized_prompts = clip.tokenize(prompts)
         with torch.inference_mode():
-            image_features = self.model.encode_image(preprocessed_image)
-            text_features = self.model.encode_text(tokenized_prompts)
+            image_features = self._model.encode_image(preprocessed_image)
+            text_features = self._model.encode_text(tokenized_prompts)
 
             # normalized features
             image_features = image_features / image_features.norm(dim=1, keepdim=True)
             text_features = text_features / text_features.norm(dim=1, keepdim=True)
 
             # cosine similarity as logits
-            logit_scale = self.model.logit_scale.exp()
+            logit_scale = self._model.logit_scale.exp()
             logits_per_image = logit_scale * image_features @ text_features.t()
 
             probs = list(logits_per_image.softmax(dim=-1).cpu().numpy()[0])
 
-        scored_prompts = {tag: float(prob) for tag, prob in zip(prompts, probs)}
-        return scored_prompts
+        return probs
     
-    def compute_images_probabilities(self, images: list[Image], prompt: str) -> dict[Image, float]:
-        raise
-        preprocessed_images = [self.img_preprocess(image).unsqueeze(0) for image in images]
+    def compute_images_probabilities(self, images: list[Image], prompt: str) -> list[float]:
+        preprocessed_images = [self._img_preprocess(image).unsqueeze(0) for image in images]
         tokenized_prompts = clip.tokenize(prompt)
         with torch.inference_mode():
-            image_features = self.model.encode_image(preprocessed_image)
-            text_features = self.model.encode_text(tokenized_prompts)
+            image_features = self._model.encode_image(torch.cat(preprocessed_images))
+            text_features = self._model.encode_text(tokenized_prompts)
 
             # normalized features
             image_features = image_features / image_features.norm(dim=1, keepdim=True)
             text_features = text_features / text_features.norm(dim=1, keepdim=True)
 
             # cosine similarity as logits
-            logit_scale = self.model.logit_scale.exp()
+            logit_scale = self._model.logit_scale.exp()
             logits_per_image = logit_scale * image_features @ text_features.t()
 
             probs = list(logits_per_image.softmax(dim=-1).cpu().numpy()[0])
 
-        scored_prompts = {tag: float(prob) for tag, prob in zip(prompts, probs)}
-        return scored_prompts
+        return probs

requirements.txt

@@ -1,4 +1,5 @@
-streamlit~=0.76.0
+streamlit~=1.11.1
 git+https://github.com/openai/CLIP@b46f5ac
 Pillow==8.1.0
-mock==4.0.3
+mock==4.0.3
+protobuf==3.20.0  # It raises errors otherwise