Update app.py

app.py

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