updates

app.py

@@ -28,7 +28,8 @@ gridsize=31
 n_patches=961
 
 #pred=np.load('pred.npy')
-pred_all=np.load('pred_all.npy').reshape(-1,64)
+pred_dict=[np.load('pred_all_hue.npy').reshape(-1,64),np.load('pred_all_grey.npy').reshape(-1,64)]
+#pred_all=np.load('pred_all.npy').reshape(-1,64)
 
 random_i=np.load('random.npy')
 
@@ -42,6 +43,10 @@ if "img" not in st.session_state:
 if "draw" not in st.session_state:
   st.session_state["draw"] = True
 
+if "model" not in st.session_state:
+  st.session_state["model"] = pred_dict.keys()[0]
+
+
 def patch(ij):
   #st.write(ij)
   immg=ij//n_patches
@@ -65,7 +70,7 @@ def find():
   #st.write(pred_all[st.session_state["img"],point[0]*36+point[1]])
   i=st.session_state["img"]
   p=point[1]*gridsize+point[0]
-  diff=np.linalg.norm(pred_all[np.newaxis,i*n_patches+p,:]-pred_all,axis=-1)
+  diff=np.linalg.norm(pred_dict[st.session_state["model"]][np.newaxis,i*n_patches+p,:]-pred_dict[st.session_state["model"]],axis=-1)
   #re_pred=pred_all.reshape(20,20,256,64)
   #diff_re=diff.reshape((20,20,256)).argmin(axis=[])
   i=0
@@ -144,6 +149,12 @@ with col1:
   with scol1:
     st.button('Change Image', on_click=button_click)
 
+    st.session_state["model"] = st.selectbox(
+    'Model',
+    pred_dict.keys())
+
+    st.write('You selected:', option)
+
   with scol2: 
     st.button('Find similar parts', on_click=find)