Update app.py

app.py

@@ -9,7 +9,16 @@ from tensorflow.keras.models import load_model
 model = load_model("unet_model.h5", compile=False)
 
 # Load Hugging Face dataset
-dataset = load_dataset("AIOmarRehan/Cropped_Yale_Faces") 
+dataset = None
+
+def get_dataset():
+    global dataset
+    if dataset is None:
+        dataset = load_dataset(
+            "AIOmarRehan/Cropped_Yale_Faces",
+            trust_remote_code=True
+        )
+    return dataset
 
 # Preprocess function
 def preprocess_image(image, target_size=(192, 176)):
@@ -38,9 +47,10 @@ def add_salt_and_pepper_noise(image, amount=0.05):
 
 # Prediction function
 def predict(img=None, use_dataset=False, add_noise=False):
+
     if use_dataset:
-        # Pick random image from dataset
-        example = random.choice(dataset["train"])
+        ds = get_dataset()
+        example = random.choice(ds["train"])
         img = example["image"]
 
     if img is None: