Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
The Raspberry Pi code running my mom's Christmas present.
import numpy
import picamera
import picamera.array
import random
import RPi.GPIO as GPIO
from time import sleep
camera_width = 320
camera_height = 240
arma_pin = 32
armb_pin = 33
trigger_pin = 40
arma_start_pw = 0
armb_start_pw = 0
def setup():
"""Configure GPIO pins and camera, returning a tuple of armA, armB, and camera."""
# GPIO PWM mapping:
# GPIO12 = PWM0 = Pin 32
# GPIO18 = PWM0 = Pin 12
# GPIO13 = PWM1 = Pin 33
# GPIO19 = PWM1 = Pin 35
GPIO.setwarnings(False)
GPIO.setmode(GPIO.BOARD) # Easier to use pin numbers IMHO.
GPIO.setup(arma_pin, GPIO.OUT)
GPIO.setup(armb_pin, GPIO.OUT)
GPIO.setup(trigger_pin, GPIO.IN, pull_up_down=GPIO.PUD_DOWN)
arma_pwm = GPIO.PWM(arma_pin, 2000)
armb_pwm = GPIO.PWM(armb_pin, 2000)
arma_pwm.start(arma_start_pw)
armb_pwm.start(armb_start_pw)
camera = picamera.PiCamera() # .rotation, .resolution, .framerate
return (arma_pwm, armb_pwm, camera)
def capture_image_to_disk(camera):
camera.start_preview()
sleep(5)
#camera.start_recording("/tmp/video")
camera.capture("/tmp/image.jpg")
camera.stop_preview()
def capture_image_as_numpy_array(camera):
output = numpy.empty((camera_height, camera_width, 3), dtype=numpy.uint8)
with picamera.array.PiRGBArray(camera) as output:
camera.resolution = (camera_width, camera_height)
camera.capture(output, 'rgb')
return (output.array.max(axis=2)//3)
def save_matrix_as_pgm(filename, matrix):
# PGM is greyscale only.
width = matrix.shape[1]
height = matrix.shape[0]
with open(filename, 'wt') as fout:
fout.write("P2\n")
fout.write(f"{width} {height}\n")
fout.write(str(int(matrix.max()))+"\n")
for y in range(height):
for x in range(width):
fout.write(str(int(matrix[y,x])) + " ")
fout.write("\n")
def matrix_to_point_array(matrix):
# Convert from this matrix into a point list.
# The darker the value, the more likely there is to be a point and the less it moves.
points = list()
width = matrix.shape[1]
height = matrix.shape[0]
mean = matrix.mean()
std = matrix.std()
high = matrix.max()
window_size = 2
for y in range(window_size, height-window_size):
for x in range(window_size, width-window_size):
# Select a block around this point and check if luminance is less than the mean.
luminance = matrix[y,x]
window = matrix[y-window_size:y+window_size, x-window_size:x+window_size]
window_mean = window.mean()
add_point = False
if luminance > mean or luminance >= window_mean or abs(luminance - window_mean) < 0.1:
add_point = False # Region is too bright or too boring.
elif luminance < window_mean or luminance < mean-2*std:
add_point = random.random() > (luminance/high) # Region is dark or interesting.
if add_point:
dx = random.random()*window_size
dy = random.random()*window_size
points.append((x+dx, y+dy))
return points
def get_path_through_points(points):
"""Given a list of points, return an array of indices which corresponds to a path through them."""
# Greedy nearest-first solution.
distance = lambda a,b: (a[0]-b[0])**2 + (a[1]-b[1])**2
visited_points = set()
unvisited_points = [i for i in range(len(points))]
path = list()
current_point = unvisited_points[0]
path.append(current_point)
unvisited_points.remove(current_point)
while unvisited_points:
print(100.0*(1.0 - len(unvisited_points)/float(len(points))))
# Find the nearest distance.
nearest_idx = unvisited_points[0]
nearest_dist = distance(points[current_point], points[nearest_idx])
for candidate_idx in unvisited_points[1:]:
dist = distance(points[current_point], points[candidate_idx])
if dist < nearest_dist:
nearest_idx = candidate_idx
nearest_dist = dist
# Have a new nearest point.
path.append(nearest_idx)
current_point = nearest_idx
unvisited_points.remove(nearest_idx)
return path
def DEBUG_get_path_through_points(points):
"""Given a list of points, return an array of indices which corresponds to a path through them."""
# Least-cost Hamiltonian path is NP-Hard. TSP is NP-Complete. Dijkstra and A* don't guarantee a path.
# Instead, we take a variant of Prim's algorithm.
neighbor_count = dict() # Keep track of the cardinality of each node.
not_in_tree = list()
in_tree = list()
edges = list()
distance = lambda a,b: (a[0]-b[0])**2 + (a[1]-b[1])**2
# We will want to find those with order 3+ in the future.
for p in points:
neighbor_count[p] = 0 # THIS WILL NOT WORK IF p IS NOT A TUPLE! It must be hashable.
not_in_tree.append(p)
# We build a minimum spanning tree.
# PRIM'S ALGORITHM
current_point = random.choice(points)
not_in_tree.remove(current_point)
in_tree.append(current_point)
while not_in_tree:
# Pick the smallest edge that connects something in the tree to something not in the tree.
new_edge_length = 1e1000
new_edge = ((0,0), (0, 0))
for p1 in in_tree:
for p2 in not_in_tree:
edge_length = distance(p1, p2)
if edge_length < new_edge_length:
new_edge_length = edge_length
new_edge = (p1, p2)
# P2 is not in the tree, so that's the one we worry about.
not_in_tree.remove(new_edge[1])
in_tree.append(new_edge[1])
edges.append(new_edge)
neighbor_count[new_edge[0]] += 1
neighbor_count[new_edge[1]] += 1
# Now we have a minimum spanning tree.
def debug_point_array_to_matrix(points):
matrix = 255 * numpy.ones((camera_height, camera_width), dtype=numpy.uint8)
for p in points:
x, y = p
matrix[int(min(matrix.shape[0]-1, max(0, y))), int(min(matrix.shape[1]-1, max(0, x)))] = 0
return matrix
def debug_point_list_to_svg(filename, points, path):
with open(filename, 'wt') as fout:
fout.write(f"<svg height=\"{camera_height}\" width=\"{camera_width}\">\n")
for path_a, path_b in zip(path, path[1:]):
p1 = points[path_a]
p2 = points[path_b]
fout.write(f"<line x1=\"{p1[0]}\" y1=\"{p1[1]}\" x2=\"{p2[0]}\" y2=\"{p2[1]}\" style=\"stroke:rgb(0, 0, 0);stroke-width:1\" />\n")
fout.write("</svg>")
def main():
pwmA, pwmB, camera = setup()
while True:
#if GPIO.input(trigger_pin):
if True:
print("Capturing...")
res = capture_image_as_numpy_array(camera)
res = (res - res.min())/(1e-6+res.max()-res.min())
print(res.shape)
save_matrix_as_pgm("/home/pi/test.pgm", res*255)
print("Pointilizing...")
points = matrix_to_point_array(res)
if not points:
print("Too few points! Image is too uniform!")
print("Pathing...")
path = get_path_through_points(points)
print("Drawing...")
#res = debug_point_array_to_matrix(points)
#save_matrix_as_pgm("/home/pi/text.pgm", res)
debug_point_list_to_svg("/home/pi/test.svg", points, path)
print("Done")
break
sleep(0.01)
if __name__=="__main__":
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.