Update app.py

app.py

@@ -3,14 +3,18 @@ import tensorflow as tf
 import keras
 from keras.datasets import mnist
 import matplotlib.pyplot as plt
+import random
 
 (train_images, train_labels), (test_images, test_labels) = mnist.load_data()
 
 def get_digit(digit_choice):
-  digit = train_images[digit_choice]
+  rn = 0
+  while(test_labels[rn] != digit_choice):
+    rn = int(random.random() * 1000)
+  digit = train_images[rn]
   fig = plt.figure()
   plt.imshow(digit)
-  out_txt = "digit: %d" % digit_choice
+  out_txt = "digit index: %d\ndigit label: %d" % digit_choice, rn
   return fig, out_txt
 
 iface = gr.Interface(