Update app.py

app.py

@@ -72,7 +72,7 @@ urllib.request.urlretrieve(MODEL_URL, "popd.pkl")
 path = Path(".")
 learn=load_learner(path, 'popd.pkl')
 
-def predict(image):
+def predict(image,colour):
   img_fast = open_image(image)
   a = PIL.Image.open(image).convert('RGB')
   st.image(a, caption='Input')
@@ -108,7 +108,7 @@ def predict(image):
       stateMap[i][j] = True
       color = int((bitmap[i, j][0] + bitmap[i, j][1] + bitmap[i, j][2])/3)
       if color > 100:
-        bitmap[i, j] =(185, 39, 40)
+        bitmap[i, j] =colour
         neigh = getNeighbours(i, j, n, m)
         for ne in neigh:
           queue.append(ne)
@@ -124,17 +124,25 @@ SIDEBAR_OPTIONS = [SIDEBAR_OPTION_DEMO_IMAGE, SIDEBAR_OPTION_UPLOAD_IMAGE]
 
 app_mode = st.sidebar.selectbox("Please select from the following", SIDEBAR_OPTIONS)
 photos = ["fight.jpg","shaolin-kung-fu.jpg"]
-
+colour = ['Red','Blue','Yellow']
 if app_mode == SIDEBAR_OPTION_DEMO_IMAGE:
   st.sidebar.write(" ------ ")
   option = st.sidebar.selectbox('Please select a sample image, then click Magic Time button', photos)
+  colour = st.sidebar.selectbox("Colour", colour)
+  colour[0] = (185, 39, 40)
+  colour[1] = (40, 96, 219)
+  colour[2] = (249, 223, 2)
   pressed = st.sidebar.button('PoP')
   if pressed:
      st.empty()
      st.sidebar.write('Please wait for the magic to happen! This may take up to a minute.')
-     predict(option)
+     predict(option,colour)
 
 elif app_mode == SIDEBAR_OPTION_UPLOAD_IMAGE:
   uploaded_file = st.file_uploader("Choose an image...")
   if uploaded_file is not None:
-    predict(uploaded_file)
+    colour = st.sidebar.selectbox("Colour", colour)
+    colour[0] = (185, 39, 40)
+    colour[1] = (40, 96, 219)
+    colour[2] = (249, 223, 2)
+    predict(uploaded_file,colour)