Update app.py

app.py

@@ -9,12 +9,12 @@ import random
 
 def pick_digit(digit_choice):
   rn = 0
-  # pick a random digit from training set until a match is found
+  # pick a random digit from 60,000 in the training set until a desired match is found
   while(train_labels[rn] != digit_choice):
-    rn = int(random.random() * 1000)  
+    rn = int(random.random() * 60000)  
   digit = train_images[rn]
   fig = plt.figure()
-  plt.imshow(digit)
+  plt.imshow(digit, cmap=plt.cm.binary)
   out_txt = "train_images[%d] = %d" % (rn, train_labels[rn])
   return fig, out_txt