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 numpy as np
from glob import glob
from collections.abc import Iterable
import matplotlib.pyplot as plt
from pathlib import Path
from matplotlib import cm
OUTPUT_GRAPH = False
OUTPUT_FRAMES = False
OUTPUT_HEIGHT_MAPS = False
X_OFFSET = 200
Y_OFFSET = 20
FRAME_SIZE_X = 200
FRAME_SIZE_Y = 60
from processing import *
from visualization import graph_height_map
from analysis import compute_x_value
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:
laser_x_val = compute_x_value(line)
frame_height_data[index] = laser_x_val
return frame_height_data
@ -50,159 +42,17 @@ def generate_height_data_from_video(video_file: str):
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):
sum_of_scores = 0
for line in height_map.transpose():
sum_of_scores += compute_score_for_frame(line)
sum_of_scores += np.std(line)
return sum_of_scores
def main():
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)
# print(frame_index, frame_std)
@ -222,7 +72,7 @@ def main():
# return
fig.suptitle(video_file)
# fig.suptitle(video_file)
# 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