Checking in before a major refactor

This commit is contained in:
Mike Abbott 2023-02-28 15:01:34 -07:00
parent 9e55c7c901
commit cf3f20a0c1

143
main.py
View File

@ -5,6 +5,8 @@ from collections.abc import Iterable
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
from pathlib import Path from pathlib import Path
OUTPUT_GRAPH = False
OUTPUT_FRAMES = False
def brightest_average(pixel_values: np.ndarray): def brightest_average(pixel_values: np.ndarray):
brightest_pixels = np.argsort(pixel_values)[-3:] brightest_pixels = np.argsort(pixel_values)[-3:]
@ -14,14 +16,15 @@ def brightest_average(pixel_values: np.ndarray):
def weighted_average(pixel_values: np.ndarray): def weighted_average(pixel_values: np.ndarray):
normalized_values = pixel_values / 255 normalized_values = pixel_values / 255
adjusted_values = normalized_values adjusted_values = normalized_values ** 200
x_values = np.arange(pixel_values.size) x_values = np.arange(adjusted_values.size)
return np.average(x_values, weights=pixel_values) return np.average(x_values, weights=adjusted_values)
def first_non_zero(pixel_values: np.ndarray): def first_non_zero(pixel_values: np.ndarray):
return np.nonzero(pixel_values)[0][0] return np.nonzero(pixel_values)[0][0]
def count_non_zero(pixel_values: np.ndarray): def count_non_zero(pixel_values: np.ndarray):
return np.count_nonzero(pixel_values) return np.count_nonzero(pixel_values)
@ -33,35 +36,119 @@ def compute_x_value(pixel_values: np.ndarray):
"first_non_zero": first_non_zero, "first_non_zero": first_non_zero,
"count_non_zero": count_non_zero, "count_non_zero": count_non_zero,
} }
return algorithms["brightest_avg"](pixel_values) # return algorithms["brightest_avg"](pixel_values)
# return algorithms["count_non_zero"](pixel_values) # return algorithms["count_non_zero"](pixel_values)
# return algorithms["first_non_zero"](pixel_values) # return algorithms["first_non_zero"](pixel_values)
# return algorithms["weighted_avg"](pixel_values) return algorithms["count_non_zero"](pixel_values)
fig = plt.figure()
from matplotlib.animation import FFMpegWriter
writer = FFMpegWriter(fps=30)
plt.ylim([0, 200])
l = None
def graph_frame(pixel_values: np.ndarray, output_file: str): def graph_frame(pixel_values: np.ndarray, output_file: str):
# fig.
return return
plt.figure() # plt.figure()
plt.plot(pixel_values) global l
plt.ylim([0, 200]) if l is None:
plt.savefig(output_file) l, = plt.plot(pixel_values)
plt.close() else:
x = np.arange(len(pixel_values))
l.set_data(x, pixel_values)
# writer.grab_frame()
# plt.savefig(output_file)
# plt.close()
return
def crop_frame(frame):
mid_y = 720//2 + 15
mid_x = 1280//2 + 30
frame = frame[mid_y-50:mid_y+50, mid_x+100:mid_x+300]
return frame
def preprocess_frame(frame):
lowerb = np.array([0, 0, 120])
upperb = np.array([255, 255, 255])
red_line = cv2.inRange(frame, lowerb, upperb)
masked_video = cv2.bitwise_and(frame,frame,mask = red_line)
gray = cv2.cvtColor(masked_video, cv2.COLOR_BGR2GRAY)
return gray
def apply_gaussian_blur(frame):
frame = cv2.GaussianBlur(frame, (3, 3), 0)
frame = cv2.GaussianBlur(frame, (3, 3), 0)
frame = cv2.GaussianBlur(frame, (11, 11), 0)
frame = cv2.GaussianBlur(frame, (11, 11), 0)
return frame
def compute_score_for_frame(x_values: Iterable): def compute_score_for_frame(x_values: Iterable):
return np.std(x_values) return np.std(x_values)
def compute_height_map(video_file):
video_data = cv2.VideoCapture(video_file)
frames = []
while video_data.isOpened():
ret, frame = video_data.read()
if not ret:
break
frame = crop_frame(frame)
frame = preprocess_frame(frame)
# frame = apply_gaussian_blur(frame)
laser_x_values = []
for line in frame:
if line.max() > 0:
laser_x_val = compute_x_value(line)
laser_x_values.append(laser_x_val)
frames.append(laser_x_values)
return frames
def graph_height_map(frames):
fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
points = []
for y, line_data in enumerate(frames):
for x, z in enumerate(line_data):
points.append(
(x, y, z)
)
x, y, z = zip(*points)
x, y, z = np.array(x), np.array(y), np.array(z)
ax.scatter(x, y, z)
fig.savefig("surface_map.png")
def main(): def main():
ranking = [] ranking = []
# i = 0
for video_file in sorted(glob("sample_data2/*")): for video_file in sorted(glob("sample_data2/*")):
# if i < 6:
# i += 1
# continue
# height_data = compute_height_map(video_file)
# graph_height_map(height_data)
# return
fig.suptitle(video_file)
video_data = cv2.VideoCapture(video_file) video_data = cv2.VideoCapture(video_file)
out = cv2.VideoWriter("out.avi", cv2.VideoWriter_fourcc('M','J','P','G'), 30, (400,400)) # out = cv2.VideoWriter("out.avi", cv2.VideoWriter_fourcc('M','J','P','G'), 30, (400,400))
mid_y = 720//2 + 15
mid_x = 1280//2 + 30
frame_index = 0 frame_index = 0
@ -70,36 +157,25 @@ def main():
ret, frame = video_data.read() ret, frame = video_data.read()
if not ret: if not ret:
break break
frame = frame[mid_y-50:mid_y+50, mid_x+100:mid_x+300]
lowerb = np.array([0, 0, 120]) frame = crop_frame(frame)
upperb = np.array([255, 255, 255]) frame = preprocess_frame(frame)
red_line = cv2.inRange(frame, lowerb, upperb) # frame = apply_gaussian_blur(frame)
masked_video = cv2.bitwise_and(frame,frame,mask = red_line)
# cv2.imwrite("test.png", masked_video)
# exit()
# out.write(masked_video)
gray = cv2.cvtColor(masked_video, cv2.COLOR_BGR2GRAY)
gray = cv2.GaussianBlur(gray, (3, 3), 0)
gray = cv2.GaussianBlur(gray, (3, 3), 0)
gray = cv2.GaussianBlur(gray, (11, 11), 0)
gray = cv2.GaussianBlur(gray, (11, 11), 0)
laser_x_values = [] laser_x_values = []
for line in gray: for line in frame:
# find the 4 brightest pixels
if line.max() > 0: if line.max() > 0:
laser_x_val = compute_x_value(line) laser_x_val = compute_x_value(line)
laser_x_values.append(laser_x_val) laser_x_values.append(laser_x_val)
graph_frame(laser_x_values, f"graphs/{Path(video_file).stem}-{frame_index}.png") if OUTPUT_GRAPH:
graph_frame(laser_x_values, f"graphs/{Path(video_file).stem}-{frame_index}.png")
# gray = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR) # gray = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
# out.write(gray) # out.write(gray)
cv2.imwrite(f"frame_data/{Path(video_file).stem}-{frame_index}.png", gray) if OUTPUT_FRAMES:
cv2.imwrite(f"frame_data/{Path(video_file).stem}-{frame_index}.png", frame)
frame_score = compute_score_for_frame(laser_x_values) frame_score = compute_score_for_frame(laser_x_values)
# print(frame_index, frame_std) # print(frame_index, frame_std)
@ -112,6 +188,7 @@ def main():
print(np.std(video_std)) print(np.std(video_std))
ranking.append((video_file, np.std(video_std))) ranking.append((video_file, np.std(video_std)))
# return
print('\nSCORES\n') print('\nSCORES\n')