Ejemplo n.º 1
0
def find_pieces(Im, board_corners, w, sq1_is_black=True, p=20):
	"""
	Given an (unwarped) 3-channel color image, finds the pieces.

	Required arguments:

		Im -- A 3-channels RGB image

		board_corners -- a list of board_corners from find_board

		w -- the classifier trained and returned by train_detector

	Keyword arguments:

		sq1_is_black -- Boolean value, true if the top left square in
			board_corners is black. Use find_board.is_first_square_black()
			to determine this if unknown.
			Default: True

		p -- an exponential factor which describes how much the board's perspective
			should be taken into account when classifying squares as pieces/non-pieces.
			Default: 20

	Return value: An 8x8 numpy array, where the (i,j)-th entry is:
		
		1 if the square contains a white piece
		-1 if the square contains a black piece
		0 if the square contains no piece
	"""

	# align the image
	Im_aligned, H = find_board.align_image(Im, board_corners)

	a1 = sqrt(sum((array(board_corners[0]) - array(board_corners[8]))**2))
	a2 = sqrt(sum((array(board_corners[72]) - array(board_corners[80]))**2))
	b1 = sqrt(sum((array(board_corners[8]) - array(board_corners[80]))**2))
	b2 = sqrt(sum((array(board_corners[0]) - array(board_corners[72]))**2))
	
        perspective = (b1/b2)**p, 1, (b1/b2)**p * (a1/a2)**p, (a1/a2)**p

	classification = _identify_occupied_squares(Im_aligned, w, perspective = perspective)

	pieces = sign(_identify_piece_color(Im_aligned, classification, sq1_is_black))

	return pieces
Ejemplo n.º 2
0
def train_detector(frames, piece_locations, board_corners, p=20):
	"""
	Given a list of frames and known piece locations, trains a piece detector.

	Required arguments:

		frames -- a python list of numpy arrays, where each array is a color image

		piece_locations -- a python list of lists of tuples describing the position
			of pieces on the board. The tuple (i,j) describes that there is a piece
			in the ith row, jth column, where the 0th row/column is the square 
			whose upper left corner is first in board_corners.

			For a frame where the pieces are in the
			default starting locations, this is just:

				[[(i,j) for in range(8) for j in range(8)]]

		board_corners -- a list of board_corners as generated by find_board
		
	Keyword arguments:

		p -- An exponential factor which describes how much the board's perspective
			should be taken into account when classifying squares as pieces/non-pieces.
			Default: 20

	Return value: train_detector returns a sklearn object describing a trained linear SVM.
	"""

	# shape an array for features and a list for classes
	features = zeros((0,N_FEATURES))
	classes = []

	# compute the perspective descriptors
	a1 = sqrt(sum((array(board_corners[0]) - array(board_corners[8]))**2))
	a2 = sqrt(sum((array(board_corners[72]) - array(board_corners[80]))**2))
	b1 = sqrt(sum((array(board_corners[8]) - array(board_corners[80]))**2))
	b2 = sqrt(sum((array(board_corners[0]) - array(board_corners[72]))**2))
	
        perspective = (b1/b2)**p, 1, (b1/b2)**p * (a1/a2)**p, (a1/a2)**p

	# for every frame, get the image patches, compute features, and add to the training data
	for k,f in enumerate(frames):
		squares = [(i,j) for i in range(8) for j in range(8)]
		
		for (i,j) in squares:
 	
			f_warp = find_board.align_image(f, board_corners)[0]
				
			# get the patch
			W = find_board.get_patch(f_warp, (i,j))

			# now compute some features
			current_features = _make_features(W,perspective)

			# add to training data
			features = vstack((features, current_features))

			# determine the class of this patch and add to training data
			if (i,j) in piece_locations[k]:
				classes.append(1)
			else:
				classes.append(-1)

	# now we have our features and classes, and so we train up a linear least squares model

	w = svm.SVC()
	w.fit(features, classes)

	return w
Ejemplo n.º 3
0
def test():

	print "Training a classifier..."

	# structures to hold frames and labels
	frames = []
	piece_locations = []
	
	# video to train on
	vid = cv2.VideoCapture('./data/full_game_Ruy_Lopez.avi')	

	# capture a frame and align it
	Im = vid.read()[1]

	board_corners = find_board.find_board(Im)	
	frames.append(Im)
	
	# these are not the correct locations, but roll with it...
	white_pieces = [(i,j) for i in range(8) for j in range(2)]
	black_pieces = [(i,j) for i in range(8) for j in range(6,8)]

	"""
	# now selectively correct the pieces
	white_pieces.remove((6,0))	
	white_pieces.remove((5,1))	
	white_pieces.append((5,2))
	white_pieces.append((5,3))

	black_pieces.remove((1,7))
	black_pieces.remove((3,6))
	black_pieces.append((2,5))
	black_pieces.append((3,3))
	"""

	piece_locations.append(white_pieces+black_pieces)

	w = train_detector(frames, piece_locations, board_corners)

	# now test it
	
	# video to test on
	vid2 = cv2.VideoCapture('./data/full_game_Ruy_Lopez.avi')	
	
	for frame in range(2400,100000,200):
		print frame
		vid2.set(cv2.cv.CV_CAP_PROP_POS_FRAMES, frame)

		# capture a frame and align it
		Im2 = vid2.read()[1]

		Im2_warp = find_board.align_image(Im2, board_corners)[0]
		Im2_circ = Im2_warp.copy()

		pieces = find_pieces(Im2, board_corners, w=w)

		white_pieces = zip(*nonzero(pieces == 1))
		black_pieces = zip(*nonzero(pieces == -1))

		print pieces

		for piece in white_pieces:
			cv2.circle(Im2_circ, ((piece[1]+1)*50 + 15, 15 + (1 + piece[0])*50), 6, (0,0,200))
		for piece in black_pieces:
			cv2.circle(Im2_circ, ((piece[1]+1)*50 + 15, 15 + (1 + piece[0])*50), 6, (0,2000,0))

		cv2.namedWindow('test',0)
		cv2.imshow('test', Im2_circ)
		cv2.waitKey(0)