Use streamlit-session-memo for session-specific caching

whitphx · Apr 27, 2023 · 5f0db2f · 5f0db2f
1 parent 6488586
commit 5f0db2f
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 10 deletions.
diff --git a/pages/1_object_detection.py b/pages/1_object_detection.py
@@ -12,6 +12,7 @@
 import cv2
 import numpy as np
 import streamlit as st
+from streamlit_session_memo import st_session_memo
 from streamlit_webrtc import WebRtcMode, webrtc_streamer
 
 from sample_utils.download import download_file
@@ -71,13 +72,12 @@ def generate_label_colors():
 download_file(PROTOTXT_URL, PROTOTXT_LOCAL_PATH, expected_size=29353)
 
 
-# Session-specific caching
-cache_key = "object_detection_dnn"
-if cache_key in st.session_state:
-    net = st.session_state[cache_key]
-else:
-    net = cv2.dnn.readNetFromCaffe(str(PROTOTXT_LOCAL_PATH), str(MODEL_LOCAL_PATH))
-    st.session_state[cache_key] = net
+@st_session_memo
+def get_model():
+    return cv2.dnn.readNetFromCaffe(str(PROTOTXT_LOCAL_PATH), str(MODEL_LOCAL_PATH))
+
+
+net = get_model()
 
 score_threshold = st.slider("Score threshold", 0.0, 1.0, 0.5, 0.05)
 

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -38,6 +38,7 @@ pandas = [
 ]
 streamlit = "^1.13.0"
 twilio = "^8.1.0"
+streamlit-session-memo = "^0.3.1"
 
 [tool.black]
 line-length = 88