From fb1a2d68c09d78c1fb679c2c681c949d674399c5 Mon Sep 17 00:00:00 2001
From: Jens Ahrensfeld <jens@jayfield.org>
Date: Tue, 9 Jul 2024 19:04:40 +0200
Subject: [PATCH] - added plot - added CLI switches for csv und plot

---
 ocv_corner_tracker.py | 47 +++++++++++++++++++++++++++++++++----------
 1 file changed, 36 insertions(+), 11 deletions(-)

diff --git a/ocv_corner_tracker.py b/ocv_corner_tracker.py
index 0f90d7a..0330083 100644
--- a/ocv_corner_tracker.py
+++ b/ocv_corner_tracker.py
@@ -3,6 +3,7 @@ import numpy as np
 import argparse
 import sys
 import csv
+import matplotlib.pyplot as plt
 
 from util import image_crop, to_rect, bbox_add_position, bbox_center, bbox_round
 
@@ -302,6 +303,8 @@ if __name__ == '__main__':
 	parser.add_argument("--track", action="store_true")
 	parser.add_argument("--path", action="store_true")
 	parser.add_argument("--loop", action="store_true")
+	parser.add_argument("--csv", action="store_true")
+	parser.add_argument("--plot", action="store_true")
 	parser.add_argument("--name", default="Default")
 	args = parser.parse_args()
 	video = cv2.VideoCapture(args.filename)
@@ -339,7 +342,8 @@ if __name__ == '__main__':
 	key_wait = -1
 	dist_min = [(0, 0)]*len(tracker_list)
 	dist_max = [(0, 0)]*len(tracker_list)
-	result_data = [[] for i in range(len(tracker_list))]
+	result_csv = [[] for i in range(len(tracker_list))]
+	result_plot = [[] for i in range(len(tracker_list))]
 	while True:
 		# Start timer
 		timer = cv2.getTickCount()
@@ -358,9 +362,12 @@ if __name__ == '__main__':
 		for ct in tracker_list:
 			mean_distance = ct.process(image, image_anno)
 			if mean_distance is not None:
-				result_data[i].append({f'Distance': mean_distance})
 				scale = ct_params.scale
 				scaled_distance = (scale*mean_distance[0], scale*mean_distance[1])
+				if args.csv:
+					result_csv[i].append({f'dx': scaled_distance[0], 'dy': scaled_distance[1]})
+				if args.plot:
+					result_plot[i].append(scaled_distance)
 				dist_min[i] = (min(scaled_distance[0], dist_min[i][0]), min(scaled_distance[1], dist_min[i][1]))
 				dist_max[i] = (max(scaled_distance[0], dist_max[i][0]), max(scaled_distance[1], dist_max[i][1]))
 				cv2.putText(image_anno, f"Distance [{i}] : ({scaled_distance[0]:+05.2f}, "
@@ -388,15 +395,33 @@ if __name__ == '__main__':
 				key_wait = -1
 
 	# writing to csv file
-	i = 1
-	for result in result_data:
-		with open(f"{args.name}_{i}.csv", 'w') as csvfile:
-			# creating a csv dict writer object
-			writer = csv.DictWriter(csvfile, fieldnames=['Distance'])
+	if len(result_csv) > 0:
+		i = 1
+		for result in result_csv:
+			filename = f"{args.name}_{i}.csv"
+			with open(filename, 'w') as csvfile:
+				# creating a csv dict writer object
+				writer = csv.DictWriter(csvfile, fieldnames=['dx', 'dy'])
 
-			# writing headers (field names)
-			writer.writeheader()
+				# writing headers (field names)
+				writer.writeheader()
 
-			# writing data rows
-			writer.writerows(result)
+				# writing data rows
+				writer.writerows(result)
+				print(f"Exported {filename}")
+				i += 1
+
+	if len(result_plot[0]) > 0:
+		N_p = len(result_plot)
+		i = 1
+		for result in result_plot:
+			Np = len
+			n = range(0, len(result))
+			plt.subplot(N_p, 1, i)
+			plt.plot(n, np.array(result)[:, 0], n, np.array(result)[:, 1])
+			plt.grid()
+			plt.legend([f"x", f"y"])
 			i += 1
+
+	plt.show()
+