From d18f9d31625ebd2244eebd376dcd613a0d93b207 Mon Sep 17 00:00:00 2001
From: Epicalert <epicalert@protonmail.com>
Date: Fri, 2 Jul 2021 02:56:39 +0800
Subject: [PATCH] Change model selection to dropdown menu

---
 CMakeLists.txt     |  1 +
 src/fc2dconfig.cpp | 25 ++++++++++++++++++++-----
 src/modellist.cpp  | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 src/modellist.hpp  |  8 ++++++++
 4 files changed, 75 insertions(+), 5 deletions(-)
 create mode 100644 src/modellist.cpp
 create mode 100644 src/modellist.hpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index dcb0eb2..933e81e 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -72,5 +72,6 @@ target_link_libraries( fc2d ${OpenCV_LIBS} ${OPENGL_LIBRARIES} ${WEBP_LIBRARIES}
 add_executable( fc2dconfig
 	src/fc2dconfig.cpp
 	src/paths.cpp
+	src/modellist.cpp
 )
 target_link_libraries( fc2dconfig  ${wxWidgets_LIBRARIES} )
diff --git a/src/fc2dconfig.cpp b/src/fc2dconfig.cpp
index c3ad3ab..199109b 100644
--- a/src/fc2dconfig.cpp
+++ b/src/fc2dconfig.cpp
@@ -3,6 +3,7 @@
 #include <config.hpp>
 #include <args.hpp>
 #include <paths.hpp>
+#include <modellist.hpp>
 
 #include <wx/wxprec.h>
 #ifndef WX_PRECOMP
@@ -18,8 +19,9 @@ class ConfigurationFrame : public wxFrame {
 	public:
 		ConfigurationFrame();
 	private:
+		std::vector<std::string> modelVec;
 		wxCheckBox* useHaarCheckBox;
-		wxTextCtrl* modelNameText;
+		wxChoice* modelNameChoice;
 		void OnExit(wxCommandEvent& event);
 };
 
@@ -41,15 +43,24 @@ ConfigurationFrame::ConfigurationFrame() : wxFrame(NULL, wxID_ANY, "Configure "
 	// TODO: load config file and populate values
 	wxStaticText* placeholderText = new wxStaticText(panel, wxID_ANY, "There's nothing here right now, just click OK.");
 	useHaarCheckBox = new wxCheckBox(panel, wxID_ANY, "Disable DNN face detection");
-	// TODO: make this a dropdown with the detected model files
-	modelNameText = new wxTextCtrl(panel, wxID_ANY, "test");
+
+	// find models to populate model dropdown
+	modelVec = listModels();
+	wxString modelArray[modelVec.size()];
+
+	for (int i = 0; i < modelVec.size(); i++) {
+		modelArray[i] = wxString(modelVec[i]);
+	}
+
+	modelNameChoice = new wxChoice(panel, wxID_ANY, wxDefaultPosition, wxDefaultSize, modelVec.size(), modelArray);
+
 	// TODO: cancel button to exit without saving settings
 	wxButton* okButton = new wxButton(panel, wxID_OK, "OK");
 
 	wxBoxSizer* sizer = new wxBoxSizer(wxVERTICAL);
 	sizer->Add(placeholderText, 0, wxALL | wxALIGN_LEFT, 10);
 	sizer->Add(useHaarCheckBox, 0, wxALL | wxALIGN_LEFT, 5);
-	sizer->Add(modelNameText, 0, wxALL | wxALIGN_LEFT, 5);
+	sizer->Add(modelNameChoice, 0, wxALL | wxALIGN_LEFT, 5);
 	sizer->AddStretchSpacer(1);
 	sizer->Add(okButton, 0, wxALIGN_RIGHT | wxALL, 10);
 
@@ -63,7 +74,11 @@ void ConfigurationFrame::OnExit(wxCommandEvent& event) {
 	std::ofstream configFile;
 	configFile.open(prefixCustom + "config.toml");
 	configFile << "use_haar = " << (useHaarCheckBox->GetValue() ? "true" : "false") << std::endl;
-	configFile << "model = \"" << modelNameText->GetValue() << "\""<< std::endl;
+	// janky edge case lmao
+	// if user did not select any model, don't write the line in the config file!
+	if (modelNameChoice->GetSelection() != wxNOT_FOUND) {
+		configFile << "model = \"" << modelVec[modelNameChoice->GetSelection()] << "\""<< std::endl;
+	}
 	configFile.close();
 
 	Close(true);
diff --git a/src/modellist.cpp b/src/modellist.cpp
new file mode 100644
index 0000000..59bc692
--- /dev/null
+++ b/src/modellist.cpp
@@ -0,0 +1,46 @@
+#include <iostream>
+
+#include <dirent.h>
+
+#include <paths.hpp>
+
+#include <modellist.hpp>
+
+// get models from a given directory and add them to a given vector
+void getModelsFromDir(std::string path, std::vector<std::string>* vector) {
+	DIR* dir;
+	struct dirent* ent;
+
+	if ((dir = opendir(path.c_str())) != NULL) {
+		while ((ent = readdir(dir)) != NULL) {
+			std::string filename = ent->d_name;
+			size_t fmaPos = filename.find(".fma"); // position of ".fma" in filename
+
+			if (fmaPos == std::string::npos) {
+				// filename does not have ".fma" in it
+				continue;
+			}
+
+			filename.erase(fmaPos);
+			
+			vector->push_back(filename);
+		}
+		closedir(dir);
+	}
+}
+
+std::vector<std::string> listModels() {
+	std::vector<std::string> modelList;
+
+	getModelsFromDir(prefixCustom + "models", &modelList);
+	getModelsFromDir(prefixDefault + "models", &modelList);
+	getModelsFromDir("models", &modelList);
+
+	/*
+	for (int i = 0; i < modelList.size(); i++) {
+		std::cout << "Detected model: " << modelList.at(i) << std::endl;
+	}
+	*/
+
+	return modelList;
+}
diff --git a/src/modellist.hpp b/src/modellist.hpp
new file mode 100644
index 0000000..ec46d39
--- /dev/null
+++ b/src/modellist.hpp
@@ -0,0 +1,8 @@
+#ifndef MODELLIST_HPP
+#define MODELLIST_HPP
+
+#include <vector>
+
+std::vector<std::string> listModels();
+
+#endif