From c18a1a9ce4607fc3ca2a92dba2d0056fc6e7df3c Mon Sep 17 00:00:00 2001
From: Roy Shilkrot <roy.shil@gmail.com>
Date: Wed, 29 May 2024 14:40:41 -0400
Subject: [PATCH] feat: Update showUnseenObjects flag and add
 saveDetectionsPath field to FilterData struct

---
 buildspec.json        |  2 +-
 data/locale/en-US.ini |  1 +
 src/FilterData.h      |  1 +
 src/detect-filter.cpp | 33 ++++++++++++++++++++++++++++++++-
 4 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/buildspec.json b/buildspec.json
index b9f361b..0735951 100644
--- a/buildspec.json
+++ b/buildspec.json
@@ -38,7 +38,7 @@
     },
     "name": "obs-detect",
     "displayName": "OBS Object Detection plugin",
-    "version": "0.0.2",
+    "version": "0.0.3",
     "author": "Roy Shilkrot",
     "website": "https://github.com/occ-ai/obs-detect",
     "email": "roy.shil@gmail.com",
diff --git a/data/locale/en-US.ini b/data/locale/en-US.ini
index 24b30a6..501c0d2 100644
--- a/data/locale/en-US.ini
+++ b/data/locale/en-US.ini
@@ -34,3 +34,4 @@ MaxUnseenFrames="Max Unseen Frames"
 ExternalModel="External Model"
 ModelPath="Model Path"
 ShowUnseenObjects="Show Currently Undetected Objects"
+SaveDetectionsPath="Save Detections Path"
diff --git a/src/FilterData.h b/src/FilterData.h
index c73c409..869c041 100644
--- a/src/FilterData.h
+++ b/src/FilterData.h
@@ -31,6 +31,7 @@ struct filter_data {
 	int lastDetectedObjectId;
 	bool sortTracking;
 	bool showUnseenObjects;
+	std::string saveDetectionsPath;
 
 	// create SORT tracker
 	Sort tracker;
diff --git a/src/detect-filter.cpp b/src/detect-filter.cpp
index aa189f6..e2f893d 100644
--- a/src/detect-filter.cpp
+++ b/src/detect-filter.cpp
@@ -53,7 +53,8 @@ static bool enable_advanced_settings(obs_properties_t *ppts, obs_property_t *p,
 	const bool enabled = obs_data_get_bool(settings, "advanced");
 
 	for (const char *prop_name :
-	     {"threshold", "useGPU", "numThreads", "model_size", "detected_object"}) {
+	     {"threshold", "useGPU", "numThreads", "model_size", "detected_object", "sort_tracking",
+	      "max_unseen_frames", "show_unseen_objects", "save_detections_path"}) {
 		p = obs_properties_get(ppts, prop_name);
 		obs_property_set_visible(p, enabled);
 	}
@@ -266,6 +267,11 @@ obs_properties_t *detect_filter_properties(void *data)
 	// add option to show unseen objects
 	obs_properties_add_bool(props, "show_unseen_objects", obs_module_text("ShowUnseenObjects"));
 
+	// add file path for saving detections
+	obs_properties_add_path(props, "save_detections_path",
+				obs_module_text("SaveDetectionsPath"), OBS_PATH_FILE_SAVE,
+				"JSON file (*.json);;All files (*.*)", nullptr);
+
 	/* GPU, CPU and performance Props */
 	obs_property_t *p_use_gpu =
 		obs_properties_add_list(props, "useGPU", obs_module_text("InferenceDevice"),
@@ -381,6 +387,7 @@ void detect_filter_defaults(obs_data_t *settings)
 	obs_data_set_default_double(settings, "zoom_factor", 0.0);
 	obs_data_set_default_double(settings, "zoom_speed_factor", 0.05);
 	obs_data_set_default_string(settings, "zoom_object", "single");
+	obs_data_set_default_string(settings, "save_detections_path", "");
 }
 
 void detect_filter_update(void *data, obs_data_t *settings)
@@ -408,6 +415,7 @@ void detect_filter_update(void *data, obs_data_t *settings)
 		tf->tracker.setMaxUnseenFrames(maxUnseenFrames);
 	}
 	tf->showUnseenObjects = obs_data_get_bool(settings, "show_unseen_objects");
+	tf->saveDetectionsPath = obs_data_get_string(settings, "save_detections_path");
 
 	// check if tracking state has changed
 	if (tf->trackingEnabled != newTrackingEnabled) {
@@ -779,6 +787,29 @@ void detect_filter_video_tick(void *data, float seconds)
 			      objects.end());
 	}
 
+	if (!tf->saveDetectionsPath.empty()) {
+		std::ofstream detectionsFile(tf->saveDetectionsPath);
+		if (detectionsFile.is_open()) {
+			nlohmann::json j;
+			for (const edgeyolo_cpp::Object &obj : objects) {
+				nlohmann::json obj_json;
+				obj_json["label"] = obj.label;
+				obj_json["confidence"] = obj.prob;
+				obj_json["rect"] = {{"x", obj.rect.x},
+						    {"y", obj.rect.y},
+						    {"width", obj.rect.width},
+						    {"height", obj.rect.height}};
+				obj_json["id"] = obj.id;
+				j.push_back(obj_json);
+			}
+			detectionsFile << j.dump(4);
+			detectionsFile.close();
+		} else {
+			obs_log(LOG_ERROR, "Failed to open file for writing detections: %s",
+				tf->saveDetectionsPath.c_str());
+		}
+	}
+
 	if (tf->preview || tf->maskingEnabled) {
 		if (tf->preview && objects.size() > 0) {
 			edgeyolo_cpp::utils::draw_objects(frame, objects, tf->classNames);