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)
@@ -117,22 +117,33 @@ def predict(image):
  
 SIDEBAR_OPTION_DEMO_IMAGE = "Select a Demo Image"
 SIDEBAR_OPTION_UPLOAD_IMAGE = "Upload an Image"
-  
-SIDEBAR_OPTIONS = [SIDEBAR_OPTION_DEMO_IMAGE, SIDEBAR_OPTION_UPLOAD_IMAGE]
+SIDEBAR_OPTION_COLOUR_IMAGE = "Choose a colour"
 
+SIDEBAR_OPTIONS = [SIDEBAR_OPTION_DEMO_IMAGE, SIDEBAR_OPTION_UPLOAD_IMAGE,SIDEBAR_OPTION_COLOUR_IMAGE]
+
+clr = [(185, 39, 40),(40, 96, 219),(249, 223, 2)]
+colour = ["Red","Blue","Yellow"]
+option2 = st.sidebar.selectbox('Please select a sample image, then click Magic Time button', colour)
+colour[0] = clr[0]
+colour[1] = clr[1]
+colour[2] = clr[2]
 
 app_mode = st.sidebar.selectbox("Please select from the following", SIDEBAR_OPTIONS)
 photos = ["fight.jpg","shaolin-kung-fu.jpg"]
-if app_mode == SIDEBAR_OPTION_DEMO_IMAGE:
+
+if app_mode == SIDEBAR_OPTION_DEMO_IMAGE and app_mode == SIDEBAR_OPTION_COLOUR_IMAGE:
   st.sidebar.write(" ------ ")
   option = st.sidebar.selectbox('Please select a sample image, then click Magic Time button', photos)
   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,option2)
 
-elif app_mode == SIDEBAR_OPTION_UPLOAD_IMAGE:
+elif app_mode == SIDEBAR_OPTION_UPLOAD_IMAGE and app_mode == SIDEBAR_OPTION_COLOUR_IMAGE:
+  colour[0] = clr[0]
+  colour[1] = clr[1]
+  colour[2] = clr[2]
   uploaded_file = st.file_uploader("Choose an image...")
   if uploaded_file is not None:
-    predict(uploaded_file)
+    predict(uploaded_file,option2)