diff --git a/nanosoc_board/C_Cxx/fast_knn_demo/fast_knn_demo.cpp b/nanosoc_board/C_Cxx/fast_knn_demo/fast_knn_demo.cpp
index e4b11c8727fd5b89d024d0a31879b96580ffa7c4..87a9d941ca5c7485c8e2a3a5ea21112055c7fd08 100644
--- a/nanosoc_board/C_Cxx/fast_knn_demo/fast_knn_demo.cpp
+++ b/nanosoc_board/C_Cxx/fast_knn_demo/fast_knn_demo.cpp
@@ -40,6 +40,7 @@ PicoGraphics_PenRGB332 graphics(st7789.width, st7789.height, nullptr);
 Pen BG = graphics.create_pen(0, 0, 0);
 Pen WHITE = graphics.create_pen(255, 255, 255);
 Pen RED = graphics.create_pen(255, 0, 0);
+Pen GREEN = graphics.create_pen(0, 255, 0);
 bool BUTTON_Y_pressed = false;
 bool BUTTON_X_pressed = false;
 int x=0;
@@ -56,18 +57,23 @@ std::string console[CONSOLE_FIFO_DEPTH];
 #define DATA_SENT_GPIO_PIN 3
 
 #define NUM_UNLABELLED_IMGS 36
-#define NUM_LABELLED_IMGS 130
+#define NUM_LABELLED_IMGS 150
 
 #define NUM_OF_CLASSES 10
 #define NUM_LABELLED_IMAGES_PER_CLASS  NUM_LABELLED_IMGS/NUM_OF_CLASSES
 #define kNN_k 5
 
+// String arrays that hold the names of image data files
 std::string labelled_files[NUM_OF_CLASSES][NUM_LABELLED_IMAGES_PER_CLASS];
 std::string unlabelled_files[NUM_UNLABELLED_IMGS];
 
+// Image data buffers
 uint8_t unlabelled_buffer[NUM_UNLABELLED_IMGS][784];
 uint8_t labelled_buffer[NUM_LABELLED_IMGS][784];
+
+// Image label buffers
 uint8_t labelled_imgs_labels[NUM_LABELLED_IMGS];
+uint8_t test_imgs_labels[NUM_UNLABELLED_IMGS];
 
 // Nanosoc image buffer address
 int nanosoc_img_buffer_addr;
@@ -84,19 +90,19 @@ uintptr_t sw_reset_reg_addr = FAST_KNN_REGS_BASE + offsetof(FAST_KNN_regs_typede
 
 // Path to labelled images 
 std::string labelled_image_dirs[10] = {
-    "./data/labelled_images/image0",
-    "./data/labelled_images/image1",
-    "./data/labelled_images/image2",
-    "./data/labelled_images/image3",
-    "./data/labelled_images/image4",
-    "./data/labelled_images/image5",
-    "./data/labelled_images/image6",
-    "./data/labelled_images/image7",
-    "./data/labelled_images/image8",
-    "./data/labelled_images/image9"};
+    "./fmnist_data/train/0_t-shirt_top",
+    "./fmnist_data/train/1_trouser",
+    "./fmnist_data/train/2_pullover",
+    "./fmnist_data/train/3_dress",
+    "./fmnist_data/train/4_coat",
+    "./fmnist_data/train/5_sandal",
+    "./fmnist_data/train/6_shirt",
+    "./fmnist_data/train/7_sneaker",
+    "./fmnist_data/train/8_bag",
+    "./fmnist_data/train/9_ankle_boot"};
 
 // Path to unlabelled images
-std::string unlabelled_image_dir = "0:/data/unlabelled_images";
+std::string unlabelled_image_dir = "0:/fmnist_data/test";
 
 std::string classes[] = {
     "T-Shirt",
@@ -344,8 +350,12 @@ void DEMO_ROUTINE(){
             result_location.x = img_x_grid[xtext];
             result_location.y = img_y_grid[ytext]+60;
 
-            // KNN finished, plot the result
-            graphics.set_pen(WHITE);
+            // KNN finished, plot the result    
+            if (predicted_label == test_imgs_labels[current_unlab_img-1])
+                graphics.set_pen(GREEN);
+            else
+                graphics.set_pen(RED);
+
             graphics.text(classes[predicted_label], result_location, result_location.x+56);
             xtext++;
             if(xtext>=3){
@@ -650,7 +660,10 @@ int main() {
         char fil_buf[8]="";
         char * pEnd;
         int k=0;
-        printf("File name: %s, Unlabelled img %u\n",  unlabelled_files[i].c_str(), i);
+
+        // Parse ground truth label of test image
+        sscanf(unlabelled_files[i].c_str(), "%hhu_*.txt", &test_imgs_labels[i]);
+        printf("File name: %s, Unlabelled img %u, label: %hhu\n",  unlabelled_files[i].c_str(), i, test_imgs_labels[i]);
         while(f_gets(fil_buf, 8, &fil)){
             unlabelled_buffer[i][k] = strtol(fil_buf,&pEnd,16);
             // printf("0x%02x ", unlabelled_buffer[j][k]);
@@ -669,9 +682,9 @@ int main() {
 
     // Download program to nanosoc
     f_chdir("./programs");
-    fr = f_open(&fil, "test_gpio.hex", FA_READ);
+    fr = f_open(&fil, "fast_knn.hex", FA_READ);
     if (FR_OK != fr && FR_EXIST != fr){
-       printf("f_open(test_gpio.hex) error: %s (%d)\n", FRESULT_str(fr), fr);
+       printf("f_open(fast_knn.hex) error: %s (%d)\n", FRESULT_str(fr), fr);
        printf("Can't open demo program. Is SD card inserted?\n");
        return 1;
     }