added image caching for hopefully easier use

app.py

@@ -31,20 +31,20 @@ st.sidebar.markdown(
 
 with st.spinner("Loading Model"):
     virtexModel, imageLoader, sample_images, valid_subs = create_objects()
+    
+    
+select_idx = None
 
-random_image = get_rand_img(sample_images)
-rand_idx = 0
 
 st.sidebar.title("Select a sample image")
 
 if st.sidebar.button("Random Sample Image"):
-    rand_idx, _ = get_rand_img(sample_images)
-    sample_image = None
+    select_idx = get_rand_idx(sample_images)
     
 sample_image = st.sidebar.selectbox(
     "",
     sample_images,
-    index=rand_idx
+    index = 0 if select_idx is None else select_idx
 )
 
 st.sidebar.title("Select a Subreddit")
@@ -66,6 +66,7 @@ with st.sidebar.form("file-uploader-form", clear_on_submit=True):
     submitted = st.form_submit_button("Submit")
     if uploaded_file is not None and submitted:
         uploaded_image = Image.open(io.BytesIO(uploaded_file.getvalue()))
+        select_idx = None # set this to help rewrite the cache
         
 
 
@@ -75,9 +76,15 @@ if uploaded_image is None and submitted:
 else:
     image_file = sample_image
     
-    image = uploaded_image if uploaded_image is not None else Image.open(image_file)
+    # LOAD AND CACHE THE IMAGE
+    if uploaded_image is not None:
+        image = uploaded_image
+    elif select_idx is None:
+        image = st.session_state.image
+    else:
+        image = Image.open(image_file)
     
-    #cache the image
+    st.session_state.image = image
     
     image_dict = imageLoader.transform(image)
     

model.py

@@ -124,9 +124,8 @@ def download_files():
 def get_samples():
     return glob.glob(SAMPLES_PATH)
 
-def get_rand_img(samples):
-    i = random.randint(0,len(samples)-1)
-    return i, samples[i]
+def get_rand_idx(samples):
+    return random.randint(0,len(samples)-1)
 
 @st.cache
 def create_objects():