Both classification and prompt ranking working

streamlit_app.py

@@ -1,3 +1,5 @@
+from typing import Optional, List
+
 from PIL import Image
 import streamlit as st
 import booste
@@ -9,76 +11,117 @@ from session_state import SessionState, get_state
 BOOSTE_API_KEY = "3818ba84-3526-4029-9dc8-ef3038697ea2"
 
 
-task_name: str = st.sidebar.radio("Task", options=["Image classification", "Image ranking", "Prompt ranking"])
+class Sections:
+    @staticmethod
+    def header():
+        st.markdown("# CLIP playground")
+        st.markdown("### Try OpenAI's CLIP model in your browser")
+        st.markdown(" ");
+        st.markdown(" ")
+        with st.beta_expander("What is CLIP?"):
+            st.markdown("Nice CLIP explaination")
+        st.markdown(" ");
+        st.markdown(" ")
+
+    @staticmethod
+    def image_uploader(accept_multiple_files: bool) -> Optional[List[str]]:
+        uploaded_image = st.file_uploader("Upload image", type=[".png", ".jpg", ".jpeg"],
+                                          accept_multiple_files=accept_multiple_files)
+
+    @staticmethod
+    def image_picker(state: SessionState):
+        col1, col2, col3 = st.beta_columns(3)
+        with col1:
+            default_image_1 = "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.image = default_image_1
+        with col2:
+            default_image_2 = "https://cdn.pixabay.com/photo/2019/12/17/18/20/peacock-4702197_960_720.jpg"
+            st.image(default_image_2, use_column_width=True)
+            if st.button("Select image 2"):
+                state.image = default_image_2
+        with col3:
+            default_image_3 = "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.image = default_image_3
+
+    @staticmethod
+    def prompts_input(state: SessionState, input_label: str, prompt_prefix: str = ''):
+        raw_classes = st.text_input(input_label)
+        if raw_classes:
+            state.prompts = [prompt_prefix + class_name for class_name in raw_classes.split(";") if len(class_name) > 1]
+            state.prompt_prefix = prompt_prefix
+
+    @staticmethod
+    def input_preview(state: SessionState):
+        col1, col2 = st.beta_columns([2, 1])
+        with col1:
+            st.markdown("Image to classify")
+            if state.image is not None:
+                st.image(state.image, use_column_width=True)
+            else:
+                st.warning("Select an image")
+
+        with col2:
+            st.markdown("Labels to choose from")
+            if state.processed_classes is not None:
+                for prompt in state.prompts:
+                    st.write(prompt[len(state.prompt_prefix):])
+            else:
+                st.warning("Enter the classes to classify from")
 
-st.markdown("# CLIP playground")
-st.markdown("### Try OpenAI's CLIP model in your browser")
-st.markdown(" "); st.markdown(" ")
-with st.beta_expander("What is CLIP?"):
-    st.markdown("Nice CLIP explaination")
-st.markdown(" "); st.markdown(" ")
+    @staticmethod
+    def classification_output(state: SessionState):
+        # Possible way of customize this https://discuss.streamlit.io/t/st-button-in-a-custom-layout/2187/2
+        if st.button("Predict"):
+            with st.spinner("Predicting..."):
+                clip_response = booste.clip(BOOSTE_API_KEY,
+                                            prompts=state.prompts,
+                                            images=[state.image],
+                                            pretty_print=True)
+                st.markdown("### Results")
+                simplified_clip_results = [(prompt[len(state.prompt_prefix):],
+                                            list(results.values())[0]["probabilityRelativeToPrompts"])
+                                           for prompt, results in clip_response.items()]
+                simplified_clip_results = sorted(simplified_clip_results, key=lambda x: x[1], reverse=True)
+
+                for prompt, probability in simplified_clip_results:
+                    percentage_prob = int(probability * 100)
+                    st.markdown(
+                        f"### ![prob](https://progress-bar.dev/{percentage_prob}/?width=200) &nbsp &nbsp {prompt}")
+                st.write(clip_response)
+
+
+task_name: str = st.sidebar.radio("Task", options=["Image classification", "Image ranking", "Prompt ranking"])
+session_state = get_state()
 if task_name == "Image classification":
-    session_state = get_state()
-    uploaded_image = st.file_uploader("Upload image", type=[".png", ".jpg", ".jpeg"],
-                                      accept_multiple_files=False)
+    Sections.header()
+    Sections.image_uploader(accept_multiple_files=False)
     st.markdown("or choose one from")
-    col1, col2, col3 = st.beta_columns(3)
-    with col1:
-        default_image_1 = "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"):
-            session_state.image = default_image_1
-    with col2:
-        default_image_2 = "https://cdn.pixabay.com/photo/2019/12/17/18/20/peacock-4702197_960_720.jpg"
-        st.image(default_image_2, use_column_width=True)
-        if st.button("Select image 2"):
-            session_state.image = default_image_2
-    with col3:
-        default_image_3 = "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"):
-            session_state.image = default_image_3
-    raw_classes = st.text_input("Enter the classes to chose from separated by a comma."
-                                " (f.x. `banana, sailing boat, honesty, apple`)")
-    if raw_classes:
-        session_state.processed_classes = raw_classes.split(",")
-        input_prompts = ["A picture of a " + class_name for class_name in session_state.processed_classes]
-
-col1, col2 = st.beta_columns([2, 1])
-with col1:
-    st.markdown("Image to classify")
-    if session_state.image is not None:
-        st.image(session_state.image, use_column_width=True)
-    else:
-        st.warning("Select an image")
-
-with col2:
-    st.markdown("Classes to choose from")
-    if session_state.processed_classes is not None:
-        for class_name in session_state.processed_classes:
-            st.write(class_name)
-    else:
-        st.warning("Enter the classes to classify from")
-
-# Possible way of customize this https://discuss.streamlit.io/t/st-button-in-a-custom-layout/2187/2
-if st.button("Predict"):
-    with st.spinner("Predicting..."):
-        clip_response = booste.clip(BOOSTE_API_KEY,
-                                    prompts=input_prompts,
-                                    images=[session_state.image],
-                                    pretty_print=True)
-        st.markdown("### Results")
-        simplified_clip_results = [(prompt[len('A picture of a '):],
-                                    list(results.values())[0]["probabilityRelativeToPrompts"])
-                                   for prompt, results in clip_response.items()]
-        simplified_clip_results = sorted(simplified_clip_results, key=lambda x: x[1], reverse=True)
-        max_class_name_length = max(len(class_name) for class_name, _ in simplified_clip_results)
-
-        for prompt, probability in simplified_clip_results:
-            progress_bar = "".join([":large_blue_circle:"] * int(probability * 10) +
-                                   [":black_circle:"] * int((1 - probability) * 10))
-            st.markdown(f"### {prompt}: {progress_bar} {probability:.3f}")
-        st.write(clip_response)
+    Sections.image_picker(session_state)
+    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 ')
+    Sections.input_preview(session_state)
+    Sections.classification_output(session_state)
+elif task_name == "Prompt ranking":
+    Sections.header()
+    Sections.image_uploader(accept_multiple_files=False)
+    st.markdown("or choose one from")
+    Sections.image_picker(session_state)
+    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)
+    Sections.input_preview(session_state)
+    Sections.classification_output(session_state)
+elif task_name == "Image ranking":
+    Sections.header()
+    Sections.image_uploader(accept_multiple_files=True)
+    st.markdown("or use random dataset")
+    Sections.image_picker(session_state)
+
+
 
 session_state.sync()