mnist_cnn: More test samples, log overall score

Author: jonnor

examples/mnist_cnn/README.md

@@ -55,7 +55,7 @@ mpremote mip install https://emlearn.github.io/emlearn-micropython/builds/master
 
 ```console
 mpremote cp mnist_cnn.tmdl :
-mpremote cp -r data/ :
+mpremote cp -r test_data/ :
 mpremote run mnist_cnn_run.py
 ```
 

examples/mnist_cnn/mnist_cnn_run.py

@@ -1,11 +1,13 @@
 
+import os
 import array
-import emlearn_cnn
 import time
 import gc
 
+import emlearn_cnn
+
 MODEL = 'mnist_cnn.tmdl'
-TEST_DATA_DIR = 'data/'
+TEST_DATA_DIR = 'test_data'
 
 def argmax(arr):
     idx_max = 0
@@ -30,6 +32,27 @@ def print_2d_buffer(arr, rowstride):
         gc.collect()
         print('\n')
 
+def load_images_from_directory(path):
+    sep = '/'
+
+    for filename in os.listdir(path):
+        # TODO: support standard image formats, like .bmp/.png/.jpeg
+        if not filename.endswith('.bin'):
+            continue
+
+        # Find the label (if any). The last part, X_label.format
+        label = None
+        basename = filename.split('.')[0]
+        tok = basename.split('_')
+        if len(tok) > 2:
+            label = tok[-1]
+
+        data_path = path + sep + filename
+        with open(data_path, 'rb') as f:
+            img = array.array('B', f.read())
+
+            yield img, label
+
 def test_cnn_mnist():
 
     # load model
@@ -42,22 +65,28 @@ def test_cnn_mnist():
     probabilities = array.array('f', (-1 for _ in range(out_length)))
 
     # run on some test data
-    for class_no in range(0, 10):
-        data_path = TEST_DATA_DIR + 'mnist_example_{0:d}.bin'.format(class_no)
-        #print('open', data_path)
-        with open(data_path, 'rb') as f:
-            img = array.array('B', f.read())
-
-            print_2d_buffer(img, 28)
-
-            run_start = time.ticks_us()
-            model.run(img, probabilities)
-            out = argmax(probabilities)
-            run_duration = time.ticks_diff(time.ticks_us(), run_start) / 1000.0 # ms
-
-            print('mnist-example-check', class_no, out, class_no == out, run_duration)
+    n_correct = 0
+    n_total = 0
+    for img, label in load_images_from_directory(TEST_DATA_DIR):
+        class_no = int(label) # mnist class labels are digits
+
+        #print_2d_buffer(img, 28)
+
+        run_start = time.ticks_us()
+        model.run(img, probabilities)
+        out = argmax(probabilities)
+        run_duration = time.ticks_diff(time.ticks_us(), run_start) / 1000.0 # ms
+        correct = class_no == out
+        n_total += 1
+        if correct:
+            n_correct += 1
+
+        print('mnist-example-check', class_no, '=', out, correct, round(run_duration, 3))
 
         gc.collect()
 
+    accuracy = n_correct / n_total
+    print('mnist-example-done', n_correct, '/', n_total, round(accuracy*100, ), '%')
+
 if __name__ == '__main__':
     test_cnn_mnist()

examples/mnist_cnn/mnist_train.py

@@ -40,7 +40,9 @@ def train_mnist(h5_file, epochs=10):
     (x_orig_train, y_orig_train), (x_orig_test, y_orig_test) = mnist.load_data() 
     num_classes = 10
 
-    generate_test_files('test_data', x_orig_test, y_orig_test)
+    TEST_DATA_DIR = 'test_data'
+    generate_test_files(TEST_DATA_DIR, x_orig_test, y_orig_test)
+    print('Wrote test data to', TEST_DATA_DIR)
 
     x_train = x_orig_train
     x_test = x_orig_test
@@ -58,7 +60,7 @@ def train_mnist(h5_file, epochs=10):
 
     model.save(h5_file)
 
-def generate_test_files(out_dir, x, y):
+def generate_test_files(out_dir, x, y, samples_per_class=5):
 
     if not os.path.exists(out_dir):
         os.makedirs(out_dir)
@@ -71,15 +73,12 @@ def generate_test_files(out_dir, x, y):
     # select one per class
     for class_no in classes:
         matches = (Y_classes == class_no)
-        print('mm', matches.shape)
         x_matches = X_series[matches]
 
-        selected = x_matches.sample(n=1, random_state=1)
-        for s in selected:
-            print('ss', s.shape, s.dtype)
-            print(s)
-            out = os.path.join(out_dir, f'mnist_example_{class_no}.bin')
-            data = s.tobytes(order='C')
+        selected = x_matches.sample(n=samples_per_class, random_state=1)
+        for i, sample in enumerate(selected):
+            out = os.path.join(out_dir, f'mnist_example_{i}_{class_no}.bin')
+            data = sample.tobytes(order='C')
 
             assert len(data) == expect_bytes, (len(data), expect_bytes)
             with open(out, 'wb') as f: