app.py updates

app.py

@@ -14,6 +14,13 @@ pred=np.load('pred.npy')
 pred_all=np.load('pred_all.npy')
 keep_bool=np.load('keep_bool.npy')
 
+keep=np.where(keep_bool.flatten())[0]
+
+keep_i=np.zeros(36*36)
+
+keep_i[keep]=keep
+
+
 
 #st.set_page_config(
 #    page_title="Streamlit Image Coordinates: Image Update",
@@ -23,6 +30,8 @@ keep_bool=np.load('keep_bool.npy')
 
 #"# :dart: Streamlit Image Coordinates: Image Update"
 
+
+
 if "point" not in st.session_state:
   st.session_state["point"] = (200,200)
 
@@ -32,13 +41,27 @@ if "img" not in st.session_state:
 if "draw" not in st.session_state:
   st.session_state["draw"] = False
 
+def patch(i):
+  img=i//(36*36)
+  p=i%(36*36)
+
+  im=ds[img]['image']
+
+  x=p//36
+  y=p%36
+  imc=im.crop((x*16,y*16,(x+1)*16,(y+1)*16))
+  return imc
+
+
+
 if "sideimg" not in st.session_state:
-  st.session_state["sideimg"] = [0,1,2,3]
+  st.session_state["sideimg"] = [patch(i) for i in range(4)]
 
 def button_click():
     st.session_state["img"]=np.random.randint(100)
     st.session_state["draw"] = False
 
+
 def find():
   point=st.session_state["point"]
   point=(point[0]//16,point[1]//16)
@@ -49,7 +72,9 @@ def find():
   p=point[0]*36+point[1]
   diff=np.linalg.norm(pred_all[np.newaxis,np.newaxis,i,p]-pred_all,axis=-1)
   for i in range(4):
-    st.write(diff.flatten().argsort()[i])
+    st.session_state["sideimg"][i]=patch(diff.argsort()[i])
+
+
 
  # for i in range(4):
  #   st.session_state["sideimg"][i]+=1
@@ -105,6 +130,6 @@ with col1:
 with col2:
 
   for i in range(3):
-    st.image(np.array(ds[st.session_state["sideimg"][i]]['image'])[::4,::4,:])
+    st.image(st.session_state["sideimg"][i])