app.py updates

app.py

@@ -12,7 +12,7 @@ ds = load_dataset("Circularmachines/batch_indexing_machine_100_small_imgs", spli
 
 
 pred=np.load('pred.npy')
-pred_all=np.load('pred_all.npy')
+pred_all=np.load('pred_all.npy').reshape(-1,64)
 keep_bool=np.load('keep_bool.npy')
 
 keep=np.where(keep_bool.flatten())[0]
@@ -72,7 +72,7 @@ def find():
   #st.write(pred_all[st.session_state["img"],point[0]*36+point[1]])
   i=st.session_state["img"]
   p=point[0]*36+point[1]
-  diff=np.linalg.norm(pred_all[np.newaxis,np.newaxis,i,p,:]-pred_all,axis=-1).flatten()
+  diff=np.linalg.norm(pred_all[np.newaxis,i*36*36+p,:]-pred_all,axis=-1)
   for ix in range(4):
     st.session_state["sideimg"][ix]=patch(diff.argsort()[ix])
     #st.write(diff.argsort()[ix])