- added plot

- added CLI switches for csv und plot
This commit is contained in:
2024-07-09 19:04:40 +02:00
parent e1f097ec8d
commit fb1a2d68c0
+36 -11
View File
@@ -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()