Update app.py

app.py

@@ -7,11 +7,18 @@ from keras.models import Model
 from datasets import load_dataset
 from sklearn.cluster import KMeans
 import matplotlib.pyplot as plt
+import pickle
 
 autoencoder = load_model("autoencoder_model.keras")
 encoded_images = np.load("X_encoded_compressed.npy")
-print("Shape of encoded_images:", encoded_images.shape)
-print("Sample encoded image:", encoded_images[0])
+
+
+# Load the dataset with the split index
+with open("dataset_with_split.pkl", "rb") as f:
+    data = pickle.load(f)
+
+dataset = data['dataset']
+split = data['split']
 
 
 num_clusters = 10  # Choose the number of clusters
@@ -38,11 +45,6 @@ def find_nearest_neighbors(encoded_images, input_image, top_n=5):
     return [(index, distances[index]) for index in nearest_neighbors]
 
 def get_image(index):
-    dataset = load_dataset('eybro/images')
-    split_dataset = dataset['train'].train_test_split(test_size=0.2)
-    dataset['train'] = split_dataset['train']
-    dataset['test'] = split_dataset['test']
-    split = len(dataset["train"])
     if index < split:
         return dataset["train"][index]
     else:
@@ -80,7 +82,6 @@ def inference(image):
     print(f"top 4: {top4}")
     
     for i in top4:
-      print(i)  
       im = get_image(i)
       print(im["label"], im["timestamp"])