mirror of
https://github.com/furrysalamander/rubedo.git
synced 2025-09-26 23:29:12 +08:00
Checking in before a major refactor
This commit is contained in:
parent
9e55c7c901
commit
cf3f20a0c1
143
main.py
143
main.py
@ -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')
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user