Update app.py

app.py

@@ -11,7 +11,8 @@ from PIL import Image
 DIMS = (100, 100)
 
 # Load the trained model
-model = load_model('mnist_model.h5')
+mnist_model = load_model('mnist_model.h5')
+adv_model = load_model('adv_model.h5')
 
 # Load MNIST examples
 (x_train, y_train), (x_test, y_test) = mnist.load_data()
@@ -33,7 +34,7 @@ def preprocess_image(image):
 # Function to make predictions
 def predict(image):
     image = preprocess_image(image)
-    prediction = model.predict(image)
+    prediction = mnist_model.predict(image)
     predicted_label = np.argmax(prediction)
     return predicted_label, np.max(prediction)
 
@@ -42,7 +43,7 @@ def get_gradients(image, label):
     image = tf.convert_to_tensor(image.reshape((1, 28, 28, 1)), dtype=tf.float32)
     with tf.GradientTape() as tape:
         tape.watch(image)
-        prediction = model(image)
+        prediction = mnist_model(image)
         loss = tf.keras.losses.categorical_crossentropy([label], prediction)
     gradients = tape.gradient(loss, image)
     return gradients.numpy().reshape(28, 28)
@@ -50,18 +51,18 @@ def get_gradients(image, label):
 # Function to progressively mask image and observe changes
 def progressively_mask_image(image, steps=100, increment=5):
     image = preprocess_image(image).reshape(28, 28)
-    label = np.argmax(model.predict(image.reshape(1, 28, 28, 1)))
+    label = np.argmax(mnist_model.predict(image.reshape(1, 28, 28, 1)))
     gradients = get_gradients(image, to_categorical(label, 10))
 
     modified_image = np.copy(image)
-    original_prediction = model.predict(image.reshape(1, 28, 28, 1))
+    original_prediction = mnist_model.predict(image.reshape(1, 28, 28, 1))
     original_label = np.argmax(original_prediction)
 
     for i in range(1, steps + 1):
         threshold = np.percentile(np.abs(gradients), 100 - i * increment)
         mask = np.abs(gradients) > threshold
         modified_image[mask] = 0
-        modified_prediction = model.predict(modified_image.reshape(1, 28, 28, 1))
+        modified_prediction = mnist_model.predict(modified_image.reshape(1, 28, 28, 1))
         predicted_label = np.argmax(modified_prediction)
         if predicted_label != original_label:
             break
@@ -88,7 +89,7 @@ def fgsm_attack(image, epsilon, data_grad):
 def create_adversarial_pattern(input_image, input_label):
     with tf.GradientTape() as tape:
         tape.watch(input_image)
-        prediction = model(input_image)
+        prediction = adv_model(input_image)
         loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)(input_label, prediction)
     gradient = tape.gradient(loss, input_image)
     return gradient