Update README.md

README.md

@@ -46,31 +46,26 @@ from skimage.util import random_noise
 import numpy as np
 import matplotlib.pyplot as plt
 
-# Load the model from Hugging Face
 model_path = hf_hub_download(repo_id="BueormLLC/AID", filename="model.h5")
-autoencoder = tf.keras.models.load_model(model_path)
 
-# Function to add noise
+autoencoder = tf.keras.models.load_model(model_path, custom_objects={'mse': tf.keras.losses.MeanSquaredError()})
+
 def add_noise(img, noise_type="gaussian"):
+    img_np = img.numpy()
     if noise_type == "gaussian":
-        noisy_img = random_noise(img, mode='gaussian', var=0.1)
+        noisy_img = random_noise(img_np, mode='gaussian', var=0.1)
     elif noise_type == "salt_pepper":
-        noisy_img = random_noise(img, mode='s&p', amount=0.1)
+        noisy_img = random_noise(img_np, mode='s&p', amount=0.1)
     return np.clip(noisy_img, 0., 1.)
 
-# Example: Load an image, add noise, and denoise
 def predict_denoised_image(autoencoder, image):
-    # Preprocess image (resize to 256x256, normalize to [0, 1])
     img_resized = tf.image.resize(image, (256, 256)) / 255.0
     img_array = np.expand_dims(img_resized, axis=0)
 
-    # Add noise
     noisy_image = add_noise(img_resized)
 
-    # Denoise the image
     denoised_image = autoencoder.predict(np.expand_dims(noisy_image, axis=0))
 
-    # Plot results
     fig, ax = plt.subplots(1, 2, figsize=(10, 5))
     ax[0].imshow(noisy_image)
     ax[0].set_title("Noisy Image")
@@ -82,11 +77,9 @@ def predict_denoised_image(autoencoder, image):
 
     plt.show()
 
-# Load a test image (replace with your own image)
-test_image = tf.keras.preprocessing.image.load_img('your_image.jpg', target_size=(256, 256))
+test_image = tf.keras.preprocessing.image.load_img('image.jpg', target_size=(256, 256))
 test_image = tf.keras.preprocessing.image.img_to_array(test_image)
 
-# Run denoising prediction
 predict_denoised_image(autoencoder, test_image)
 ```