- added plot
- added CLI switches for csv und plot
This commit is contained in:
+36
-11
@@ -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()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user