Checking in the separated code

Next up is generating the PA code from Python.
This commit is contained in:
Mike Abbott 2023-03-28 21:10:54 -06:00
parent f594ecc63c
commit 7cabd9107a
5 changed files with 151 additions and 157 deletions

41
analysis.py Normal file
View File

@ -0,0 +1,41 @@
import numpy as np
from constants import *
def brightest_average(pixel_values: np.ndarray):
brightest_pixels = np.argsort(pixel_values)[-3:]
line_brightest_x = FRAME_SIZE_X - np.average(brightest_pixels)
return line_brightest_x
def weighted_average(pixel_values: np.ndarray):
normalized_values = pixel_values / 255
adjusted_values = normalized_values ** 100
x_values = np.arange(adjusted_values.size)
if adjusted_values.max() == 0:
# FIXME: I need an appropriate solution for what to do if there are no non-zero values.
return 70
return FRAME_SIZE_X - np.average(x_values, weights=adjusted_values)
def first_non_zero(pixel_values: np.ndarray):
try:
return np.nonzero(pixel_values)[0][0]
except:
print()
def count_non_zero(pixel_values: np.ndarray):
return np.count_nonzero(pixel_values)
def compute_x_value(pixel_values: np.ndarray):
algorithms = {
"brightest_avg": brightest_average,
"weighted_avg": weighted_average,
"first_non_zero": first_non_zero,
"count_non_zero": count_non_zero,
}
# return algorithms["brightest_avg"](pixel_values)
# return algorithms["count_non_zero"](pixel_values)
# return algorithms["first_non_zero"](pixel_values)
return algorithms["weighted_avg"](pixel_values)

9
constants.py Normal file
View File

@ -0,0 +1,9 @@
OUTPUT_GRAPH = False
OUTPUT_FRAMES = False
OUTPUT_HEIGHT_MAPS = True
X_OFFSET = 200
Y_OFFSET = 20
FRAME_SIZE_X = 200
FRAME_SIZE_Y = 60

164
main.py Normal file → Executable file
View File

@ -1,21 +1,12 @@
#!/usr/bin/python3
import cv2 import cv2
import numpy as np import numpy as np
from glob import glob from glob import glob
from collections.abc import Iterable
import matplotlib.pyplot as plt
from pathlib import Path from pathlib import Path
from matplotlib import cm
from processing import *
OUTPUT_GRAPH = False from visualization import graph_height_map
OUTPUT_FRAMES = False from analysis import compute_x_value
OUTPUT_HEIGHT_MAPS = False
X_OFFSET = 200
Y_OFFSET = 20
FRAME_SIZE_X = 200
FRAME_SIZE_Y = 60
def generate_height_data_for_frame(frame: np.ndarray): def generate_height_data_for_frame(frame: np.ndarray):
@ -29,6 +20,7 @@ def generate_height_data_for_frame(frame: np.ndarray):
# if line.max() > 0: # if line.max() > 0:
laser_x_val = compute_x_value(line) laser_x_val = compute_x_value(line)
frame_height_data[index] = laser_x_val frame_height_data[index] = laser_x_val
return frame_height_data return frame_height_data
@ -50,159 +42,17 @@ def generate_height_data_from_video(video_file: str):
return height_data return height_data
def brightest_average(pixel_values: np.ndarray):
brightest_pixels = np.argsort(pixel_values)[-3:]
line_brightest_x = FRAME_SIZE_X - np.average(brightest_pixels)
return line_brightest_x
def weighted_average(pixel_values: np.ndarray):
normalized_values = pixel_values / 255
adjusted_values = normalized_values ** 200
x_values = np.arange(adjusted_values.size)
return FRAME_SIZE_X - np.average(x_values, weights=adjusted_values)
def first_non_zero(pixel_values: np.ndarray):
try:
return np.nonzero(pixel_values)[0][0]
except:
print()
def count_non_zero(pixel_values: np.ndarray):
return np.count_nonzero(pixel_values)
def compute_x_value(pixel_values: np.ndarray):
algorithms = {
"brightest_avg": brightest_average,
"weighted_avg": weighted_average,
"first_non_zero": first_non_zero,
"count_non_zero": count_non_zero,
}
# return algorithms["brightest_avg"](pixel_values)
# return algorithms["count_non_zero"](pixel_values)
# return algorithms["first_non_zero"](pixel_values)
return algorithms["weighted_avg"](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):
# fig.
return
# plt.figure()
global l
if l is None:
l, = plt.plot(pixel_values)
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 + Y_OFFSET
mid_x = 1280//2 + X_OFFSET
half_y = FRAME_SIZE_Y / 2
half_x = FRAME_SIZE_X / 2
frame = frame[int(mid_y-half_y):int(mid_y+half_y), int(mid_x-half_x):int(mid_x+half_x)]
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):
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(z_data: np.ndarray, output_file: str):
fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
# fig, ax = plt.subplots()
# 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)
y = np.arange(len(z_data))
x = np.arange(len(z_data[0]))
(x ,y) = np.meshgrid(x,y)
ax.plot_surface(x, y, z_data,cmap=cm.coolwarm,linewidth=0, antialiased=False)
# ax.pcolormesh(x, y, z_data, cmap='RdBu')
# ax.scatter(x, y, z)
fig.savefig(output_file)
def compute_score_from_heightmap(height_map: np.ndarray): def compute_score_from_heightmap(height_map: np.ndarray):
sum_of_scores = 0 sum_of_scores = 0
for line in height_map.transpose(): for line in height_map.transpose():
sum_of_scores += compute_score_for_frame(line) sum_of_scores += np.std(line)
return sum_of_scores return sum_of_scores
def main(): def main():
ranking = [] ranking = []
# if OUTPUT_GRAPH:
# graph_frame(laser_x_values, f"graphs/{Path(video_file).stem}-{frame_index}.png")
# 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)
@ -222,7 +72,7 @@ def main():
# return # return
fig.suptitle(video_file) # fig.suptitle(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))

31
processing.py Normal file
View File

@ -0,0 +1,31 @@
from constants import *
import cv2
import numpy as np
def crop_frame(frame):
mid_y = 720//2 + Y_OFFSET
mid_x = 1280//2 + X_OFFSET
half_y = FRAME_SIZE_Y / 2
half_x = FRAME_SIZE_X / 2
frame = frame[int(mid_y-half_y):int(mid_y+half_y), int(mid_x-half_x):int(mid_x+half_x)]
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

63
visualization.py Normal file
View File

@ -0,0 +1,63 @@
import cv2
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from pathlib import Path
from constants import *
def graph_frame(pixel_values: np.ndarray, output_file: str):
# fig.
return
# plt.figure()
global l
if l is None:
l, = plt.plot(pixel_values)
else:
x = np.arange(len(pixel_values))
l.set_data(x, pixel_values)
# writer.grab_frame()
# plt.savefig(output_file)
# plt.close()
return
def graph_height_map(z_data: np.ndarray, output_file: str):
# fig, ax = plt.subplots(subplot_kw={"projection": "3d"})
fig, ax = plt.subplots()
# 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)
y = np.arange(len(z_data))
x = np.arange(len(z_data[0]))
(x ,y) = np.meshgrid(x,y)
# ax.plot_surface(x, y, z_data,cmap=cm.coolwarm,linewidth=0, antialiased=False)
ax.pcolormesh(x, y, z_data, cmap='RdBu')
# ax.scatter(x, y, z)
fig.savefig(output_file)
def generate_graph_from_heightmap():
if OUTPUT_GRAPH:
graph_frame(laser_x_values, f"graphs/{Path(video_file).stem}-{frame_index}.png")
pass
def generate_frames_from_heightmap():
if OUTPUT_FRAMES:
cv2.imwrite(f"frame_data/{Path(video_file).stem}-{frame_index}.png", frame)
pass
fig = plt.figure()
from matplotlib.animation import FFMpegWriter
writer = FFMpegWriter(fps=30)
plt.ylim([0, 200])
l = None