def get_plot(vid_name):
print("--------------------------------------")
print(vid_name[3:])
# Run CV for given video
comp_vis = CV(vid_name)
img = comp_vis.run_cv()
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# Estimate trajectory for the bottom of the ball
all_data_bottom = np.asarray(comp_vis.bottom_ball)
x_bottom_data = all_data_bottom[:, 0]
y_bottom_data = all_data_bottom[:, 1]
# Ordinary least squares (OLS)
bottom_w_OLS = PolyRegression(x_bottom_data, y_bottom_data, 2, OLS)
bottom_w_OLS.get_weight_vector()
# Total least squares (TLS)
bottom_w_TLS = PolyRegression(x_bottom_data, y_bottom_data, 2, TLS)
bottom_w_TLS.get_weight_vector()
# RANSAC
bottom_w_RANSAC = PolyRegression(x_bottom_data, y_bottom_data, 2, RANSAC)
bottom_w_RANSAC.get_weight_vector()
# Estimate trajectory for the bottom of the ball
all_data_top = np.asarray(comp_vis.top_ball)
x_top_data = all_data_top[:, 0]
y_top_data = all_data_top[:, 1]
# OLS
top_w_OLS = PolyRegression(x_top_data, y_top_data, 2, OLS)
top_w_OLS.get_weight_vector()
# TLS
top_w_TLS = PolyRegression(x_top_data, y_top_data, 2, TLS)
top_w_TLS.get_weight_vector()
# RANSAC
top_w_RANSAC = PolyRegression(x_top_data, y_top_data, 2, RANSAC)
top_w_RANSAC.get_weight_vector()
# Plot results
plt.imshow(img)
plt.plot(x_top_data,
top_w_OLS.predict(x_top_data),
label=f'OLS, {top_w_OLS.eqn_str}')
plt.plot(x_top_data,
top_w_TLS.predict(x_top_data),
label=f'TLS, {top_w_TLS.eqn_str}')
plt.plot(x_top_data,
top_w_RANSAC.predict(x_top_data),
label=f'RANSAC, {top_w_RANSAC.eqn_str}')
plt.legend()
plt.title(f"Top Trajectory\n{vid_name[3:]}")
plt.grid(True)
plt.figure()
plt.imshow(img)
plt.plot(x_bottom_data,
bottom_w_OLS.predict(x_bottom_data),
label=f'OLS, {bottom_w_OLS.eqn_str}')
plt.plot(x_bottom_data,
bottom_w_TLS.predict(x_bottom_data),
label=f'TLS, {bottom_w_TLS.eqn_str}')
plt.plot(x_bottom_data,
bottom_w_RANSAC.predict(x_bottom_data),
label=f'RANSAC, {bottom_w_RANSAC.eqn_str}')
plt.legend()
plt.title(f"Bottom Trajectory\n{vid_name[3:]}")
plt.grid(True)
plt.show()
print("")
