Use client server approach for model to UI communication.

.gitignore

@@ -86,4 +86,7 @@ target/
 .mypy_cache/
 
 # exclude generated models from source control
-models/intermediate/
+models/intermediate/
+
+# exclude uploaded images from source control
+images/upload_*

app.py

@@ -1,32 +1,21 @@
+import uvicorn
 import streamlit as st
+from multiprocessing import Process
 
-from model import predictor
-from streamlit.scriptrunner import add_script_run_ctx
+import socket
+import time
 import audiobot
 import chatbot
 import os
-import threading
 
 
-def runInThread():
-    print('Initialize model in thread')
-    st.session_state['predictor'] = predictor.Predictor()
-    print('Model is initialized')
-
-
-def run():
+def run_st_app():
     st.set_page_config(
         page_title='Welcome to Visual Question Answering - Bot',
         page_icon=':robot:',
         layout='wide'
     )
 
-    os.environ['TOKENIZERS_PARALLELISM'] = 'false'
-    if 'thread' not in st.session_state:
-        st.session_state.thread = threading.Thread(target=runInThread)
-        add_script_run_ctx(st.session_state.thread)
-        st.session_state.thread.start()
-
     st.sidebar.title('VQA Bot')
     st.sidebar.write('''
         VQA Bot addresses the challenge of visual question answering with the chat and voice assistance.
@@ -41,7 +30,33 @@ def run():
     elif selected_page == 'VQA Audiobot':
         audiobot.show()
 
-    st.caption("Created by Madhuri Sakhare - [Github](https://github.com/msak1612/vqa_chatbot) [Linkedin](https://www.linkedin.com/in/madhuri-sakhare/)")
+    st.caption(
+        'Created by Madhuri Sakhare - [Github](https://github.com/msak1612/vqa_chatbot) [Linkedin](https://www.linkedin.com/in/madhuri-sakhare/)')
+
+
+def run_uvicorn():
+    os.system('uvicorn server:app --port 8080 --host 0.0.0.0 --workers 1')
+
+
+def start_server():
+    if not is_port_in_use(8080):
+        with st.spinner(text='Loading models...'):
+            proc = Process(target=run_uvicorn, args=(), daemon=True)
+            proc.start()
+            while not is_port_in_use(8080):
+                time.sleep(1)
+            st.success('Models are loaded.')
+
+
+def is_port_in_use(port):
+    # Find whether port is available using https://stackoverflow.com/questions/2470971/fast-way-to-test-if-a-port-is-in-use-using-python
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        return s.connect_ex(('0.0.0.0', port)) == 0
+
 
+if __name__ == '__main__':
+    run_st_app()
+    if 'server' not in st.session_state:
+        st.session_state['server'] = True
+        start_server()
 
-run()

audiobot.py

@@ -1,4 +1,3 @@
-from turtle import width
 import streamlit as st
 from PIL import Image
 from bokeh.models.widgets import Button
@@ -7,28 +6,14 @@ from st_clickable_images import clickable_images
 from streamlit_bokeh_events import streamlit_bokeh_events
 from bokeh.models.widgets.buttons import Button
 import time
-from os.path import *
-from os import listdir
-import base64
-
-def update_gallery_images():
-    if 'gallery' not in st.session_state:
-        st.session_state.gallery = []
-        st.session_state.gallery_images = []
-        image_path = join(dirname(abspath(__file__)), 'images')
-        for f in listdir(image_path):
-            if f.startswith('image'):
-                with open(join(image_path, f), "rb") as image:
-                    encoded = base64.b64encode(image.read()).decode()
-                    st.session_state.gallery.append(
-                        f"data:image/jpeg;base64,{encoded}")
-                    st.session_state.gallery_images.append(join(image_path, f))
+from helper import *
 
 
 def upload_image_callback():
-    st.session_state.uploaded_image = st.session_state.uploader
+    st.session_state.uploaded_image = upload_image_to_server()
     st.session_state.input = ''
 
+
 def show():
     st.session_state.audio_answer = ''
 
@@ -63,8 +48,10 @@ def show():
             on_change=upload_image_callback, key='uploader')
 
         if st.session_state.uploaded_image is not None:
-            st.session_state.image = Image.open(st.session_state.uploaded_image)
-            st.image(st.session_state.uploaded_image, use_column_width='always')
+            st.session_state.image = Image.open(
+                st.session_state.uploaded_image)
+            st.image(st.session_state.uploaded_image,
+                     use_column_width='always')
         else:
             st.session_state.image = None
             st.session_state.input = ''
@@ -118,10 +105,8 @@ def show():
                 if 'question' not in st.session_state or st.session_state.question != result.get('GET_TEXT'):
                     st.session_state['question'] = result.get('GET_TEXT')
                     with st.spinner('Preparing answer...'):
-                        while 'predictor' not in st.session_state:
-                            time.sleep(2)
-                        st.session_state.audio_answer = st.session_state.predictor.predict_answer_from_text(
-                            st.session_state.image, result.get('GET_TEXT'))
+                        st.session_state.audio_answer = request_answer(
+                            st.session_state.server_image_file, result.get('GET_TEXT'))
 
         tts_button = Button(label='Get Answer', width=100)
         tts_button.js_on_event('button_click', CustomJS(code=f"""

chatbot.py

@@ -1,12 +1,9 @@
-import streamlit as st
 from streamlit_chat import message
 from st_clickable_images import clickable_images
 from PIL import Image
-import time
-from os.path import *
-from os import listdir
-import base64
+from helper import *
 
+import streamlit as st
 
 def init_chat_history():
     if 'question' not in st.session_state:
@@ -29,39 +26,22 @@ def predict(image, input):
     if image is None or not input:
         return
 
-    if 'predictor' not in st.session_state:
-        with st.spinner('Preparing answer...'):
-            while 'predictor' not in st.session_state:
-                time.sleep(2)
-
-    answer = st.session_state.predictor.predict_answer_from_text(image, input)
-    st.session_state.question.append(input)
-    st.session_state.answer.append(answer)
-    while len(st.session_state.question) >= 5:
-        st.session_state.answer.pop(0)
-        st.session_state.question.pop(0)
-
-
-def update_gallery_images():
-    if 'gallery' not in st.session_state:
-        st.session_state.gallery = []
-        st.session_state.gallery_images = []
-        image_path = join(dirname(abspath(__file__)), 'images')
-        for f in listdir(image_path):
-            if f.startswith('image'):
-                with open(join(image_path, f), "rb") as image:
-                    encoded = base64.b64encode(image.read()).decode()
-                    st.session_state.gallery.append(
-                        f"data:image/jpeg;base64,{encoded}")
-                    st.session_state.gallery_images.append(join(image_path, f))
+    with st.spinner('Preparing answer...'):
+        answer = request_answer(st.session_state.uploaded_image, input)
+        st.session_state.question.append(input)
+        st.session_state.answer.append(answer)
+        while len(st.session_state.question) >= 5:
+            st.session_state.answer.pop(0)
+            st.session_state.question.pop(0)
 
 
 def upload_image_callback():
-    st.session_state.uploaded_image = st.session_state.uploader
+    st.session_state.uploaded_image = upload_image_to_server()
     st.session_state.question = []
     st.session_state.answer = []
     st.session_state.input = ''
 
+
 def show():
     init_chat_history()
 
@@ -98,7 +78,8 @@ def show():
 
         if st.session_state.uploaded_image is not None:
             image = Image.open(st.session_state.uploaded_image)
-            st.image(st.session_state.uploaded_image, use_column_width='always')
+            st.image(st.session_state.uploaded_image,
+                     use_column_width='always')
         else:
             st.session_state.question = []
             st.session_state.answer = []

helper.py

@@ -0,0 +1,43 @@
+from os import listdir
+from os.path import *
+from PIL import Image
+from io import BytesIO
+
+import streamlit as st
+import base64
+import requests
+
+
+def update_gallery_images():
+    if 'gallery' not in st.session_state:
+        st.session_state['gallery'] = []
+        st.session_state['gallery_images'] = []
+        image_path = join(dirname(abspath(__file__)), 'images')
+        for f in listdir(image_path):
+            if f.startswith('image'):
+                with open(join(image_path, f), "rb") as image:
+                    encoded = base64.b64encode(image.read()).decode()
+                    st.session_state.gallery.append(
+                        f"data:image/jpeg;base64,{encoded}")
+                    st.session_state.gallery_images.append(join(image_path, f))
+
+
+def upload_image_to_server():
+    if st.session_state.uploader is not None:
+        image = Image.open(st.session_state.uploader)
+        byte_io = BytesIO()
+        image.save(byte_io, 'png')
+        byte_io.seek(0)
+        file = {'file': byte_io}
+        response = requests.post('http://0.0.0.0:8080/uploadfile/', files=file)
+        if response.status_code == 200:
+            return response.json()['filename']
+    return None
+
+
+def request_answer(image, question):
+    response = requests.get(
+        f'http://0.0.0.0:8080/vqa?image={image}&question={question}')
+    if response.status_code == 200:
+        return response.json()['answer']
+    return 'I do not understand. Please ask again.'

model/predictor.py

@@ -1,4 +1,3 @@
-import streamlit as st
 from happytransformer import HappyTextToText, TTSettings
 from transformers import ViltProcessor
 from transformers import ViltForQuestionAnswering
@@ -18,7 +17,6 @@ question.
 '''
 
 
-@st.experimental_singleton
 class Predictor:
     def __init__(self):
         auth_token = os.environ.get('TOKEN') or True

requirements.txt

@@ -3,6 +3,7 @@ aiosignal==1.2.0
 altair==4.2.0
 ansicolors==1.1.8
 ansiwrap==0.8.4
+anyio==3.6.1
 appnope==0.1.3
 argon2-cffi==21.3.0
 argon2-cffi-bindings==21.2.0
@@ -29,12 +30,14 @@ defusedxml==0.7.1
 dill==0.3.5.1
 entrypoints==0.4
 executing==0.8.3
+fastapi==0.78.0
 fastjsonschema==2.15.3
 filelock==3.7.1
 frozenlist==1.3.0
 fsspec==2022.5.0
 gitdb==4.0.9
 GitPython==3.1.27
+h11==0.13.0
 happytransformer==2.4.1
 huggingface-hub==0.7.0
 idna==3.3
@@ -74,6 +77,8 @@ parso==0.8.3
 pexpect==4.8.0
 pickleshare==0.7.5
 Pillow==9.1.1
+pox==0.3.1
+ppft==1.7.6.5
 prometheus-client==0.14.1
 prompt-toolkit==3.0.29
 protobuf==3.20.1
@@ -82,12 +87,14 @@ ptyprocess==0.7.0
 pure-eval==0.2.2
 pyarrow==8.0.0
 pycparser==2.21
+pydantic==1.9.1
 pydeck==0.7.1
 Pygments==2.12.0
 Pympler==1.0.1
 pyparsing==3.0.9
 pyrsistent==0.18.1
 python-dateutil==2.8.2
+python-multipart==0.0.5
 pytz==2022.1
 pytz-deprecation-shim==0.1.0.post0
 PyYAML==6.0
@@ -106,9 +113,11 @@ simplere==1.2.13
 six==1.12.0
 sklearn==0.0
 smmap==5.0.0
+sniffio==1.2.0
 soupsieve==2.3.2.post1
 st-clickable-images==0.0.3
 stack-data==0.3.0
+starlette==0.19.1
 streamlit==1.10.0
 streamlit-bokeh-events==0.1.2
 streamlit-chat==0.0.2.1
@@ -128,6 +137,7 @@ typing_extensions==4.2.0
 tzdata==2022.1
 tzlocal==4.2
 urllib3==1.26.9
+uvicorn==0.18.1
 validators==0.20.0
 wcwidth==0.2.5
 webencodings==0.5.1

server.py

@@ -0,0 +1,68 @@
+from fastapi import FastAPI, File, UploadFile
+from model import predictor
+from os import listdir
+from os.path import *
+from PIL import Image
+
+import os
+import hashlib
+import threading
+import time
+
+gpredictor = None
+app = FastAPI()
+
+@app.get('/')
+def root():
+    return {'app': 'Thanks for visiting!!'}
+
+
+@app.get('/favicon.ico', include_in_schema=False)
+@app.post('/uploadfile/')
+async def create_upload_file(file: UploadFile = File(...)):
+    contents = await file.read()
+    hash = hashlib.sha256(contents).hexdigest()
+    file.filename = f'images/upload_{hash}.jpg'
+    if not os.path.isfile(file.filename):
+        with open(file.filename, 'wb') as f:
+            f.write(contents)
+    images[file.filename] = Image.open(file.filename)
+    return {'filename': file.filename}
+
+
+@app.get('/vqa')
+async def answer(
+    image: str,
+    question: str
+):
+    if image not in images:
+        print('not in image')
+        pil_image = Image.open(image)
+        images[image] = pil_image
+    else:
+        pil_image = images[image]
+    while gpredictor is None:
+        time.sleep(1)
+    answer = gpredictor.predict_answer_from_text( pil_image, question )
+    return {'answer': answer }
+
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+images={}
+
+def runInThread():
+    collect_images()
+    print('Initialize model in thread')
+    global gpredictor
+    gpredictor = predictor.Predictor()
+    print('Model is initialized')
+
+
+def collect_images():
+    image_path = join(dirname(abspath(__file__)), 'images')
+    for f in listdir(image_path):
+        if f.startswith('image'):
+            full_image_path = join(image_path, f)
+            images[full_image_path] = Image.open(full_image_path)
+
+thread = threading.Thread(target=runInThread)
+thread.start()