diff --git a/include/ButtonThread.h b/include/ButtonThread.h
index 5d3a2d9..c2f7855 100644
--- a/include/ButtonThread.h
+++ b/include/ButtonThread.h
@@ -7,16 +7,9 @@
 #define BUTTONTHREAD_H
 
 #include "mbed.h"
-#include <stdio.h>
-
-int readButtonPresses(void);
-
-void button_onrise(void);
 
 void buttonTask(void);
 
-bool waitPressCondition(Kernel::Clock::time_point timeout);
-
 int waitForPresses(std::chrono::milliseconds wait_time);
 
 #endif // BUTTONTHREAD_H
diff --git a/include/Filesystem.h b/include/Filesystem.h
index d03d952..9b5ec30 100644
--- a/include/Filesystem.h
+++ b/include/Filesystem.h
@@ -6,9 +6,10 @@
 #ifndef FILESYSTEM_H
 #define FILESYSTEM_H
 
-#include "MainFilesystem.h"
+#include "FileSystem.h"
+#include "mbed.h"
 
-void mountFilesystem(mbed::BlockDevice &bd, LittleFileSystem2 &fs);
-void unmountFilesystem(mbed::BlockDevice &bd, LittleFileSystem2 &fs);
+void mountFilesystem(mbed::BlockDevice &bd, mbed::FileSystem &fs);
+void unmountFilesystem(mbed::BlockDevice &bd, mbed::FileSystem &fs);
 
 #endif
diff --git a/include/FlashErase.h b/include/FlashErase.h
index cab4d7c..250c777 100644
--- a/include/FlashErase.h
+++ b/include/FlashErase.h
@@ -6,9 +6,8 @@
 #ifndef FLASHERASE_H
 #define FLASHERASE_H
 
-#include "Filesystem.h"
 #include "MainBD.h"
 
-void doErase(mbed::BlockDevice &bd, LittleFileSystem2 &fs);
+void doErase(mbed::BlockDevice &bd);
 
 #endif
diff --git a/include/MyUSBMSD.h b/include/MyUSBMSD.h
index 1ee05ce..51c7241 100644
--- a/include/MyUSBMSD.h
+++ b/include/MyUSBMSD.h
@@ -19,5 +19,5 @@
 	virtual const uint8_t *string_iproduct_desc();
 };
 
-void doUSBMSD(mbed::BlockDevice &bd, LittleFileSystem2 &fs);
+void doUSBMSD(mbed::BlockDevice &bd);
 #endif
diff --git a/src/ButtonThread.cpp b/src/ButtonThread.cpp
index bc20282..8468ff6 100644
--- a/src/ButtonThread.cpp
+++ b/src/ButtonThread.cpp
@@ -9,7 +9,6 @@
 EventFlags buttonFlags;
 
 #define PUSH_FLAG 0b01
-#define LIFT_FLAG 0b10
 
 void button_onrise(void) { buttonFlags.set(PUSH_FLAG); }
 
@@ -27,7 +26,7 @@
 		if (timed_out) {
 			break;
 		} else if (had_any_error) {
-			error("Uhh what the f?\n");
+			error("Unable to wait for push button flag?\n");
 		} else {
 			pressed_count++;
 		}
diff --git a/src/Filesystem.cpp b/src/Filesystem.cpp
index 029423b..a580a81 100644
--- a/src/Filesystem.cpp
+++ b/src/Filesystem.cpp
@@ -15,7 +15,7 @@
 
 // function to attempt to mount the filesystem and
 // reformat and mount it again should it fail
-void mountFilesystem(mbed::BlockDevice &bd, LittleFileSystem2 &fs) {
+void mountFilesystem(mbed::BlockDevice &bd, mbed::FileSystem &fs) {
 	int err;
 	// Try to mount the filesystem
 	printf("Mounting the filesystem... ");
@@ -43,7 +43,7 @@
 }
 
 // function to unmount the filesystem
-void unmountFilesystem(mbed::BlockDevice &bd, LittleFileSystem2 &fs) {
+void unmountFilesystem(mbed::BlockDevice &bd, mbed::FileSystem &fs) {
 	// Try to unmount the filesystem
 	printf("Unmounting the filesystem... ");
 	fflush(stdout);
diff --git a/src/FlashErase.cpp b/src/FlashErase.cpp
index f772fd9..a84f71d 100644
--- a/src/FlashErase.cpp
+++ b/src/FlashErase.cpp
@@ -9,9 +9,7 @@
 #include <functional>
 
 // function to erase the flash
-void doErase(mbed::BlockDevice &bd, LittleFileSystem2 &fs) {
-	// Try to unmount the filesystem
-	unmountFilesystem(bd, fs);
+void doErase(mbed::BlockDevice &bd) {
 	printf("Initializing the block device... ");
 	fflush(stdout);
 	int err = bd.init();
diff --git a/src/MyUSBMSD.cpp b/src/MyUSBMSD.cpp
index 0b23a56..08d6bbc 100644
--- a/src/MyUSBMSD.cpp
+++ b/src/MyUSBMSD.cpp
@@ -5,7 +5,6 @@
 
 #include "MyUSBMSD.h"
 #include "ButtonThread.h"
-#include "Filesystem.h"
 #include "mbed.h"
 
 // create a custom usb mass storage subclass to override
@@ -26,9 +25,7 @@
 	return string_iproduct_descriptor;
 }
 
-void doUSBMSD(mbed::BlockDevice &bd, LittleFileSystem2 &fs) {
-	// Try to unmount the filesystem
-	unmountFilesystem(bd, fs);
+void doUSBMSD(mbed::BlockDevice &bd) {
 	printf("Switching to the usb mass storage mode...\n");
 	MyUSBMSD usb(&bd, true, 0x1209, 0x10);
 
@@ -40,7 +37,6 @@
 			       "device...\n");
 			usb.disconnect();
 			printf("Switched usb mass storage mode off.\n");
-			mountFilesystem(bd, fs);
 			break;
 		}
 	}
diff --git a/src/main.cpp b/src/main.cpp
index 2c834bb..008984b 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -22,15 +22,20 @@
 
 	while (true) {
 		int presses = waitForPresses(600s);
+		unmountFilesystem(mainBD, mainFS);
+
 		if (presses == 1) {
-			doUSBMSD(mainBD, mainFS);
+			doUSBMSD(mainBD);
 		}
 
 		if (presses == 2) {
-			doErase(mainBD, mainFS);
+			doErase(mainBD);
 		}
+
+		mountFilesystem(mainBD, mainFS);
 	}
 
+	unmountFilesystem(mainBD, mainFS);
 	buttonThread.terminate();
 
 	return 0;