Rubik's Cube detection written in Python

cuberead5.pu

from picamera.array import PiYUVArray
from picamera import PiCamera
import time
import cv2
import numpy as np
import colorsys
 
#set Pi Camera settings
camera = PiCamera()
camera.resolution = (640, 480)
camera.brightness = 40
bright = 40
camera.contrast = 0
camera.saturation = -10
camera.rotation = 180
camera.framerate = 10
rawCapture = PiYUVArray(camera, size=(640, 480))
# allow the camera to initilize
time.sleep(0.1)

#set the lists for all 6 sides to white for initialize
top= ("W", "W", "W", "W", "W", "W", "W", "W", "W")
bottom=("W","W","W","W","W","W","W","W","W")
left=("W","W","W","W","W","W","W","W","W")
right=("W","W","W","W","W","W","W","W","W")
front=("W","W","W","W","W","W","W","W","W")
back=("W","W","W","W","W","W","W","W","W")
#set flags that all side are not completed
top_complete= False
bottom_complete = False
left_complete = False
right_complete = False
front_complete = False
back_complete = False

#initialize working face to nil
face=[]

#routine for counting number of each color has appeaded on each face. Used at the end to ensure there is 9 of each colour before exiting
def count_colours (fcheck, countall):
 (wcount,bcount,rcount,gcount,ycount,ocount)=countall
 for pos in range (9):
  if fcheck[pos]=="W":
   wcount +=1
  elif fcheck[pos] == "B":
   bcount +=1
  elif fcheck[pos] == "R":
   rcount+=1
  elif fcheck[pos] == "G":
   gcount+=1
  elif fcheck[pos] == "Y":
   ycount+=1
  elif fcheck[pos] == "O":
   ocount+=1
 countall=(wcount,bcount,rcount,gcount,ycount,ocount)
 return countall


 #routine for drawing the read faces on the screen
def draw_face(colorss,x,y): 
 toprow = colorss[0:3]
 midrow = colorss[3:6]
 lastrrow = colorss[6:9]
 tile_color = (0,0,0)
 #draw a black square as background
 cv2.rectangle(image, (0+x,0+y ), (65+x, 65+y), (0,0,0),-1)
 
 #draw the coloured squares for top row
 for pos in range(3):
  if toprow[pos] == "W":
   tile_color = (255,255,255)
  elif toprow[pos] == "G":
   tile_color = (0,255,0)
  elif toprow[pos] == "B":
   tile_color = (255,0,0)
  elif toprow[pos] == "R":
   tile_color = (0,0,255)
  elif toprow[pos] == "O":
   tile_color = (0,100,255)
  elif toprow[pos] == "Y":
   tile_color = (50,255,255)   
  elif toprow[pos] == "N":
   tile_color = (0,0,0)
  cv2.rectangle((image), (20*pos+5+x,5+y ), (20*pos+20+x, 20+y), tile_color,-1)
 
#draw the coloured squares for middle row 
 for pos in range(3):
  if midrow[pos] == "W":
   tile_color = (255,255,255)
  elif midrow[pos] == "G":
   tile_color = (0,255,0)
  elif midrow[pos] == "B":
   tile_color = (255,0,0)
  elif midrow[pos] == "R":
   tile_color = (0,0,255)
  elif midrow[pos] == "O":
   tile_color = (0,100,255)
  elif midrow[pos] == "Y":
   tile_color = (50,255,255)
  elif middle_row[pos] == "N":
   tile_color = (0,0,0)   
  cv2.rectangle((image), (20*pos+5+x,25+y ), (20*pos+20+x, 40+y), tile_color,-1)

#draw the coloured squares for bottom row  
 for pos in range(3):
  if lastrrow[pos] == "W":
   tile_color = (255,255,255)
  elif lastrrow[pos] == "G":
   tile_color = (0,255,0)
  elif lastrrow[pos] == "B":
   tile_color = (255,0,0)
  elif lastrrow[pos] == "R":
   tile_color = (0,0,255)
  elif lastrrow[pos] == "O":
   tile_color = (0,100,255)
  elif lastrrow[pos] == "Y":
   tile_color = (50,255,255)
  elif lastrrow[pos] =="N":
   tile_color = (0,0,0)
  cv2.rectangle((image), (20*pos+5+x,45+y ), (20*pos+20+x, 60+y), tile_color,-1)
  
    
#setup videowriter to record frames, this can beroved for faster processing
fourcc = cv2.VideoWriter_fourcc(*'XVID') 
writer = None
writer = cv2.VideoWriter("/home/pi/Shane/outvideo.avi", fourcc, 5, (640, 480), True)

#loop for capturing frames from camera in raw YUV format for the purpose of adjusting brightness
for frame in camera.capture_continuous(rawCapture, format="raw", use_video_port=True):
	
	# grab the raw NumPy array representing the image, then initialize the timestamp	
	image = frame.array
	#find the mean of all the image
	(Y,U,V,DA)= cv2.mean(image)
	#clear the image ready for next capture
	rawCapture.truncate(0)
	
	#if mean brightness : Y is between 70 and 72 then exit loop
	if Y > 70 and Y < 72:
	  break
	#adjust camera brightness based upon mean Y of image
	if Y > 82:
	 bright -=3
	if Y > 72:
	 bright -=1
	 camera.brightness = bright
	if Y < 60 :
	 bright += 3
	if Y < 70 :
	 bright+=1
	 camera.brightness = bright 

 
# Main loop for capturing frames from camera in raw YUV format
for frame in camera.capture_continuous(rawCapture, format="raw", use_video_port=True):
	
	# grab the raw NumPy array representing the image, then initialize the timestamp	
	image = frame.array		
	
	#create serval biniary masks to issolate the colours we are looking for
	White_Yellow_mask = cv2.inRange(image, (100,40,90), (255,255,255))		
	Red_mask= cv2.inRange(image, (0,0,143), (255,255,255))
	Green_mask = cv2.inRange(image, (60,50,50), (150,170,110))
	Center_Green_Mask = cv2.inRange(image, (20,60,50), (80,130,110))
	Blue_Mask = cv2.inRange(image, (20,95,50), (90,200,115))	
	
	#Combined all the biniary masks together to get one image with all needed data
	Combined_image = cv2.bitwise_or(White_Yellow_mask, Red_mask)
	Combined_image = cv2.bitwise_or(Green_mask, Combined_image)
	Combined_image = cv2.bitwise_or(Center_Green_Mask, Combined_image)
	Combined_image = cv2.bitwise_or(Blue_Mask, Combined_image)	
	
    #look for countours in the combined binary mask
	im2,contours2, hierarchy = cv2.findContours(Combined_image.copy(),cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)	
	candidates=[]
	index = 0
    # loop through the cpountours to find the ones we want	
	for c in contours2:
	    #approxPolyDP to find strainght line contours
	 	peri = cv2.arcLength(c, True)
		approx = cv2.approxPolyDP(c, 0.06 * peri, True)
		#find countours made up for 4 lines
		if len (approx) ==4 :
		 (x, y, w, h) = cv2.boundingRect(approx)
		 #find aspect ratio of boundary rectangle around the countours
		 ar = w / float(h)
		 #get the area of ther countour
		 area= cv2.contourArea(contours2[index])		 
		 #if the countour has AR close to 1 i.e close to square not rectangle and area of countour is close to area of box around countouour i.e not diamonmd or robus add it to candidates list
		 if ar > .7 and ar < 1.3 and w > 30 and w < 90 and area/(w*h) > .4:		  
		  candidates.append((x,y,w,h))		  
		index += 1
	
	
	new=candidates
    #loop through all the countours in the candidantes list	
	for d in new:
	 neighbors=0
	 (x,y,w,h) = d
	 for (x2,y2,w2,h2) in new:
	  #count up how many neighbors closer than width * 3.5 each countour in candidanteslist has
	  if abs(x-x2) < (w*3.5) and abs(y-y2) < (h*3.5):
	   neighbors +=1
     #any candidantes with less than 5 neighbors remove	   
	 if neighbors < 5 :	  
	  candidates.remove(d)
	  
	#sort candidates if there's 9 of them
	tmp=[]	
	if len(candidates)==9 :
	 #Write to tmp the center y,x of candidantes so that we can sort in y direction
	 for (x3,y3,w3,h3) in candidates:	  
	  tmp.append( (y3+(h/2), x3+(w/2)) )	 
	 tmp = sorted(tmp)
	 #cut into sets of 3 i.e the 3 rows of colours
	 top_row=tmp[0:3]
	 tmp = tmp[3:9]
	 tmp = sorted(tmp)
	 middle_row = tmp[0:3]
	 bottom_row = tmp[3:6]
	 
	 #sort top_row
	 temp_row=[]
     #write to temp_row the center x,y for sorting in x direction	 
	 for (y4,x4) in top_row:
	  temp_row.append((x4,y4))
	 #sort top_row
	 top_row = temp_row
	 top_row = sorted(top_row)
	 
	 #sort middle_row
	 temp_row=[]
	 for (y4,x4) in middle_row:
	  temp_row.append((x4,y4))
	 middle_row = temp_row
	 middle_row = sorted(middle_row)
	 
	 #sort bottom_row
	 temp_row=[]
	 for (y4,x4) in bottom_row:
	  temp_row.append((x4,y4))
	 bottom_row = temp_row
	 bottom_row = sorted(bottom_row)	
	 
	 
	 face=[]
	 #loop through the 3 positions in each row for the purpose to detacting colour of each tile
	 for pos in range(3):
	  #cut out a 10x10 cube around center of contour
	  x,y = top_row[pos]
	  cube = image[y-5:y+5, x-5:x+5]
      #find the mean of that cube	  
	  (Y,U,V,DA)= cv2.mean(cube)
      #identify each color and write the detected colour to the face list	  
	  if Y > 120 and float(U/V)>0.9 :	   
	   face.append("W")
	  elif U > 130 and U > V and float (U/Y)> 1.15:
	   face.append("B")	   
	  elif float(U/V) > 1.1 and float(U/V) < 2:
	   face.append("G")
	  elif V > 120 and float (U/Y) > 0.7:
	   if float(U/Y) < 1.9 :
	    face.append("O")		
	   else :
	    face.append("R")
	  elif Y > 110 and float(V/U)>0.95  :
	   face.append("Y")

	 #just doing same as above but for middle row
	 for pos in range(3):	  
	  x,y = middle_row[pos]
	  cube = image[y-5:y+5, x-5:x+5]      
	  (Y,U,V,DA)= cv2.mean(cube)	 
	  if Y > 120 and float(U/V)>0.9 :	   
	   face.append("W")
	  elif U > 130 and U > V and float (U/Y)> 1.15:
	   face.append("B")
	  elif pos == 1 and float(U/V) > 1 and float(U/V) < 1.9:
	   face.append("G")
	  elif float(U/V) > 1.1 and float(U/V) < 1.9:
	   face.append("G")
	  elif V > 120 and float (U/Y) > 0.7:
	   if float(U/Y) < 1.9 :
	    face.append("O")		
	   else :
	    face.append("R")
	  elif Y > 110 and  float(V/U)>0.95  :
	   face.append("Y")

	 #one more time for bottom row
	 for pos in range(3):	  
	  x,y = bottom_row[pos]
	  cube = image[y-5:y+5, x-5:x+5]      
	  (Y,U,V,DA)= cv2.mean(cube)	
	  if Y > 120 and float(U/V)>0.9 :	   
	   face.append("W")	  
	  elif U > 130 and U > V and float (U/Y)> 1.15:
	   face.append("B")	   
	  elif float(U/V) > 1.1 and float(U/V) < 2:
	   face.append("G")
	  elif V > 120 and float (U/Y) > 0.7:
	   if float(U/Y) < 1.9 :
	    face.append("O")		
	   else :
	    face.append("R")
	  elif Y > 110 and float(V/U)>0.95 :
	   face.append("Y")	
	
	#convert the image from raw YUV format to RGB for user viewing
	image = cv2.cvtColor(image, cv2.COLOR_YUV2RGB)
	#if there was 9 colours dectected then check middle tile to know which face to update.
	if len(face) == 9:
	 #draw the current face in top left of viewing image
	 draw_face(face,0,0)
	 #based unpon which colour the center tile is update that face for the detected colours and switch the flag that face has been read completly
	 if face[4] == "W":
	  top = face
	  top_complete = True
	 elif face[4] == "G":
	  front = face
	  front_complete = True
	 elif face[4] == "R":
	  right = face
	  right_complete = True
	 elif face[4] == "O":
	  left = face
	  left_complete = True
	 elif face[4] == "B":
	  back = face
	  back_complete = True
	 elif face[4] == "Y":
	  bottom = face
	  bottom_complete = True
	 
	#draw the all the faces on the viewing image
	draw_face(left,0,350)
	draw_face(front, 65,350)
	draw_face(top,65,285)
	draw_face(bottom,65,415)
	draw_face(right,130,350)
	draw_face(back,195,350)

	#draw circles of all the candidates on the viewing image
	for (x,y,w,h) in  candidates:	 
	 cv2.circle(image, (x+w/2, y+h/2), int(w/1.8), (255,0,255),3)
	
	#Display the viewing image		
	cv2.imshow("Orginial", image)
	
	#write image to video file	
	writer.write(image)	
 
	# clear the stream in preparation for the next frame
	rawCapture.truncate(0)
	
	#check to see if all faces have been read
	if top_complete and bottom_complete and front_complete and back_complete and left_complete and right_complete:
	 #if all faces read then count up how many of each colour has been detacted
	 countall = (0,0,0,0,0,0)
	 countall = count_colours(top, countall)
	 countall = count_colours(bottom, countall)
	 countall = count_colours(front,countall)
	 countall = count_colours(back, countall)
	 countall = count_colours(left, countall)
	 countall = count_colours(right, countall)
	 (wcount,bcount,rcount,gcount,ycount,ocount)=countall
	 #if there is 9 of each colour then write result to file and break from loop
	 if bcount==9 and rcount==9 and gcount==9 and wcount==9 and ycount==9 and ocount==9:
	  cubefile = open('/home/pi/Sage/MacTwist.txt',"w")	  
	  cubefile.write("Front"+"\n")      
	  for ch in front:       
	   cubefile.write(ch)	  
	  cubefile.write("\n")     
	  cubefile.write("Bottom"+"\n")      
	  for ch in bottom:	   
	   cubefile.write(ch)	  
	  cubefile.write("\n")      
	  cubefile.write("Left"+"\n")	  
	  for ch in left:       
	   cubefile.write(ch)      
	  cubefile.write("\n")      
	  cubefile.write("Right"+"\n")
	  for ch in right:
	   cubefile.write(ch)
	  cubefile.write("\n")
	  cubefile.write("Top"+"\n")
	  for ch in top:
	   cubefile.write(ch)
	  cubefile.write("\n")
	  cubefile.write("Back"+"\n")
	  for ch in back:
	   cubefile.write(ch)
	  cubefile.write("\n")
	  cubefile.close()
	  break	 
	
	#Check for key pressed
	key = cv2.waitKey(1) & 0xFF 
	# if the `q` key was pressed, break from the loop
	if key == ord("q"):
		break
#close video file
writer.release()