Raspberry Pi Camera Capture of Rubik's Cube State

First we import the libraries we need related to the Camera and image and Machine Learning libraries


In [ ]:
from picamera import PiCamera
from picamera.color import Color

from gpiozero import LED

import os
from time import sleep

# import a bunch of stuff that we'll use to manipulate our images...
from skimage.io import imread
#from skimage import filter
from skimage.segmentation import slic
from skimage.segmentation import mark_boundaries
from skimage.util import img_as_float
from skimage import io
from skimage.measure import block_reduce
import numpy as np

from sklearn.cluster import KMeans
import pandas as pd

import kociemba

%matplotlib inline

import matplotlib.pyplot as plt

In [ ]:
# Import GPIO Libraries
import RPi.GPIO as GPIO
import time

In [ ]:
#GPIO.cleanup()

In [ ]:
# Set up output pins for the two servo motors
GPIO.setmode(GPIO.BCM)
GPIO.setup(18, GPIO.OUT)
GPIO.setup(19, GPIO.OUT)

In [ ]:
# A class to extend the PWM object to be a bit smarter
class Servo(GPIO.PWM):
    """ Smarter PWM interface to servos.  This adds a set_position method and keeps
        the state of the current position
    """
    def __init__(self, pin_number, freq, start_angle, begin_duty_cycle=2.5, 
                 end_duty_cycle=12.5, sweep_angle=120.0):
        super(Servo, self).__init__(pin_number, freq)
        
        self.angle = start_angle
        self.servo_slope = (sweep_angle)/(end_duty_cycle-begin_duty_cycle)
        self.servo_intercept = -(self.servo_slope)*begin_duty_cycle
        duty_cycle = self._map_angle_to_duty_cycle(start_angle)
        self.start(duty_cycle)
        
    def set_position(self, angle, steps=10, delay=0.1):
        
        for step in np.linspace(self.angle, angle, steps):
            duty_cycle = self._map_angle_to_duty_cycle(step)
            self.ChangeDutyCycle(duty_cycle)
            time.sleep(delay)
        self.angle = angle
        
    def _map_angle_to_duty_cycle(self, angle):
        return (angle-self.servo_intercept)/self.servo_slope

In [ ]:
# Some initial settings which seem to work for my build
start_angle = 22
hold_angle = 42
push_angle = 60
angle_0 = 3
angle_90 = angle_0+65
roll_map = [4, 0, 2, 1, 3, 5]
rotate_map = [0, 2, 4, 3, 5, 1]
unrotate_map = [0, 5, 1, 3, 2, 4]

In [ ]:
def roll_cube(servo_obj, orientation):
    servo_obj.set_position(push_angle, delay=0.06)
    servo_obj.set_position(start_angle)
    orientation = orientation[roll_map]
    return orientation
    
def rotate_cube(servo_obj, orientation):
    servo_obj.set_position(angle_90-7)
    orientation = orientation[rotate_map]
    return orientation
    
def unrotate_cube(servo_obj, orientation):
    servo_obj.set_position(angle_0)
    orientation = orientation[unrotate_map]
    return orientation
    
def hold_cube(servo_obj):
    servo_obj.set_position(hold_angle)
    
def return_cube(servo_obj):
    servo_obj.set_position(start_angle)
    
def rotate_side(servo_roll, servo_side, orientation):
    hold_cube(servo_roll)
    servo_side.set_position(angle_90, delay=0.1)
    return_cube(servo_roll)
    servo_side.set_position(angle_0, delay=0.06)
    orientation = orientation[unrotate_map]
    return orientation

def counter_rotate_side(servo_roll, servo_side, orientation):
    servo_side.set_position(angle_90-5.0, delay=0.06)
    hold_cube(servo_roll)
    servo_side.set_position(angle_0, delay=0.1)
    return_cube(servo_roll)
    orientation = orientation[rotate_map]
    return orientation

In [ ]:
# Initialize the two servos
servo_roll = Servo(18, 50, start_angle)
servo_side = Servo(19, 50, angle_0)

In [ ]:
# 'Full' servo range of motion for spin
servo_side.set_position(0)
servo_side.set_position(110)

In [ ]:
# A little "wake up" movement from our servos to get them settled.
servo_side.set_position(angle_0)
servo_side.set_position(angle_90)
servo_side.set_position(angle_0)
servo_roll.set_position(hold_angle)
servo_roll.set_position(start_angle)

In [ ]:
# Initialize light
led = LED(17)

In [ ]:

Begin scan of cube faces


In [ ]:
face_order = np.array(['U', 'R', 'F', 'D', 'L', 'B'])
orientation = face_order.copy()

In [ ]:
camera = PiCamera()

In [ ]:
camera.resolution = (120, 120)

camera.hflip = False
camera.vflip = False
camera.brightness = 50 # the default is 50, but you can set it to whatever.

How about some text on the image.


In [ ]:
camera.annotate_foreground = Color(1.0,1.0,0.5)

camera.annotate_text = ""
camera.annotate_text_size = 36

In [ ]:
# Turn on the LED
led.on()

In [ ]:
camera.start_preview()
sleep(1)
camera.capture('./img/cubeU.jpg')
camera.stop_preview()

In [ ]:
# Move cube to show left side with down side at base of camera
orientation = roll_cube(servo_roll, orientation)
orientation = rotate_cube(servo_side, orientation)

In [ ]:
camera.start_preview()
sleep(1)
camera.capture('./img/cubeL.jpg')
camera.stop_preview()

In [ ]:
orientation = roll_cube(servo_roll, orientation)

In [ ]:
camera.start_preview()
sleep(1)
camera.capture('./img/cubeB.jpg')
camera.stop_preview()

In [ ]:
orientation = roll_cube(servo_roll, orientation)

In [ ]:
camera.start_preview()
sleep(1)
camera.capture('./img/cubeR.jpg')
camera.stop_preview()

In [ ]:
orientation = roll_cube(servo_roll, orientation)

In [ ]:
camera.start_preview()
sleep(1)
camera.capture('./img/cubeF.jpg')
camera.stop_preview()

In [ ]:
orientation = unrotate_cube(servo_side, orientation)
orientation = roll_cube(servo_roll, orientation)
orientation = rotate_cube(servo_side, orientation)

In [ ]:
camera.start_preview()
sleep(1)
camera.capture('./img/cubeD.jpg')
camera.stop_preview()

In [ ]:
orientation = unrotate_cube(servo_side, orientation)

In [ ]:
led.off()

In [ ]:
camera.close()

In [ ]:
orientation

Now process the images...


In [ ]:


In [ ]:
face_images = {}
face_images_out = []
squares = {}
masks = {}  # for QA

In [ ]:
for face in face_order:
    # read and cache images in dict
    face_images[face] = (imread('./img/cube%s.jpg' % face))

    img = face_images[face]
    mask = np.empty(img.shape[:2], dtype=np.bool)
    mask[::]=False
    row_coords = [(25, 30), (50, 55), (75, 80)]
    col_coords = [(38, 43), (65, 70), (95, 100)]
    squares[face] = np.zeros((3, 3, 3))
    row = 0
    # extract average RGB values from approximate square centers
    for i in row_coords:
        col = 0
        for j in col_coords:
            mask[i[0]:i[1], j[0]:j[1]] = True
            squares[face][row, col] = (np.average(img[i[0]:i[1], j[0]:j[1], 0]),
                                 np.average(img[i[0]:i[1], j[0]:j[1], 1]),
                                 np.average(img[i[0]:i[1], j[0]:j[1], 2]))
            col+=1
        row+=1
    # to show last image alignment
    masks[face] = mark_boundaries(img, mask)

In [ ]:
fig_mask = plt.figure()
ax = fig_mask.add_subplot(1,1,1)
ax.imshow(masks['D'])
plt.axis("off")

In [ ]:
square_list = []
for face in face_order:
    for x in range(3):
        for y in range(3):
            r, g, b = squares[face][x,y]
            square_list.append([face, x, y, r, g, b])

columns = ['face', 'x', 'y', 'r', 'g', 'b']
square_frame = pd.DataFrame(square_list, columns=columns)

In [ ]:
square_frame.head(10)

In [ ]:
X = square_frame[['r', 'g', 'b']]

In [ ]:
km = KMeans(n_clusters=6, random_state=123)
km.fit(X)

In [ ]:
km.cluster_centers_

In [ ]:
palette = np.zeros((100,600,3))
#print list(reversed([int(x) for x in km.cluster_centers_[0]]))
palette[0:99,0:99] = list(([int(x) for x in km.cluster_centers_[0]]))
palette[0:99,100:199] = list(([int(x) for x in km.cluster_centers_[1]]))
palette[0:99,200:299] = list(([int(x) for x in km.cluster_centers_[2]]))
palette[0:99,300:399] = list(([int(x) for x in km.cluster_centers_[3]]))
palette[0:99,400:499] = list(([int(x) for x in km.cluster_centers_[4]]))
palette[0:99,500:599] = list(([int(x) for x in km.cluster_centers_[5]]))

In [ ]:
centers_list = []
center_rgb = square_frame[(square_frame['x']==1) & (square_frame['y']==1)]
print center_rgb

In [ ]:
square_mapper = {}
center_map = km.predict(center_rgb[['r', 'g', 'b']])
for i in range(6):
    square_mapper[center_map[i]] = face_order[i]

In [ ]:
center_map

In [ ]:
square_mapper

In [ ]:
square_predict = km.predict(square_frame[['r', 'g', 'b']])

In [ ]:
squares = []
for square in square_predict:
    squares.append(square_mapper[square])

In [ ]:
cube = ''.join(squares)

In [ ]:
cube

Solve it!


In [ ]:
solution = kociemba.solve(cube)

In [ ]:
len(solution.split(' '))

In [ ]:
solution

In [ ]:


In [ ]:
orientation = np.array(['D', 'F', 'R', 'U', 'B', 'L'],
      dtype='|S1')

In [ ]:
def reorient(move_string, orientation):
    for move in move_string.split(' '):
        if move == 'roll':
            orientation = roll_cube(servo_roll, orientation)
        if move == 'rotate':
            orientation = rotate_cube(servo_side, orientation)
        if move == 'unrotate':
            orientation = unrotate_cube(servo_side, orientation)
    return orientation

In [ ]:
# this dict lists moves to get the desired index into the '3' position (on bottom)
reorient_map = {0: "roll roll", 1: "roll", 2: "rotate roll unrotate",
                3: "", 4: "roll roll roll", 5: "rotate roll roll roll unrotate"}

In [ ]:
#orientation = reorient(reorient_map[5], orientation)

In [ ]:
solution_list = solution.split(' ')
for move in solution_list:
    print "premove orientation = %s" % orientation
    if move[0] != orientation[3]:  # not already on bottom
        curr_move_side = np.argwhere(orientation==move[0])[0][0]
        orientation = reorient(reorient_map[curr_move_side], orientation)
    print "reorient = %s" % orientation
    # rotate the side
    if len(move)>1:
        if move[1]=="'":
            orientation = rotate_side(servo_roll, servo_side, orientation)
            orientation = rotate_side(servo_roll, servo_side, orientation)
            orientation = rotate_side(servo_roll, servo_side, orientation)
        elif move[1]=="2":
            orientation = rotate_side(servo_roll, servo_side, orientation)
            orientation = rotate_side(servo_roll, servo_side, orientation)
    else:
        orientation = rotate_side(servo_roll, servo_side, orientation)

In [ ]:


In [ ]:
camera.close()

In [ ]: