working on a cv experiment for the pool table

pooltable.py

import cv2
import numpy as np
import operator
from numpy import pi, sin, cos
import math

## Tweak these
hue_threshold = [4,3] # the lower and upper threshold for hue values in the mask
median_kernel_size = 5
video_mode = False
video_source = 0


def create_capture(source = 0, fallback = None):
    '''source: <int> or '<int>|<filename>|synth [:<param_name>=<value> [:...]]'
    '''
    source = str(source).strip()
    chunks = source.split(':')
    # handle drive letter ('c:', ...)
    if len(chunks) > 1 and len(chunks[0]) == 1 and chunks[0].isalpha():
        chunks[1] = chunks[0] + ':' + chunks[1]
        del chunks[0]

    source = chunks[0]
    try: source = int(source)
    except ValueError: pass
    params = dict( s.split('=') for s in chunks[1:] )

    cap = None
    if source == 'synth':
        Class = classes.get(params.get('class', None), VideoSynthBase)
        try: cap = Class(**params)
        except: pass
    else:
        cap = cv2.VideoCapture(source)
        if 'size' in params:
            w, h = map(int, params['size'].split('x'))
            cap.set(cv2.CAP_PROP_FRAME_WIDTH, w)
            cap.set(cv2.CAP_PROP_FRAME_HEIGHT, h)
    if cap is None or not cap.isOpened():
        print('Warning: unable to open video source: ', source)
        if fallback is not None:
            return create_capture(fallback, None)
    return cap

def processImage(img, video_mode = False):
	############################################
	# Step 1
	# read in the image and convert it to HSV
	############################################
	hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
	height, width, channels = img.shape
	total_area = height * width

	############################################
	# Step 2
	# build a histogram of the hue with a range of 0-255 and 255 buckets
	############################################
	bins = [255]
	range = [0,255]
	histogram = cv2.calcHist([hsv],[0],None,bins,range)
	cv2.normalize(histogram,histogram,0,255,cv2.NORM_MINMAX)
	hist = np.int32(np.around(histogram))
	max_hue, value = max(enumerate(hist), key=operator.itemgetter(1))

	############################################
	# Step 3
	# create a mask of only the colors in the hue range
	############################################
	mask = cv2.inRange(hsv, np.array([max_hue-hue_threshold[0], 0, 0]), np.array([max_hue+hue_threshold[1], 255, 255]))

	############################################
	# Step 4
	# median filter the mask
	############################################
	filtered_mask = cv2.medianBlur(mask, median_kernel_size);

	############################################
	# write out some of these steps for visualization
	############################################
	if not video_mode:
		cv2.imwrite("hsv.jpg", hsv)
		cv2.imwrite("mask.jpg", mask)
		cv2.imwrite("filtered_mask.jpg", filtered_mask)

	############################################
	# Step 5
	# find the contour around the table
	############################################
	# start by finding all of the contours in the filtered mask
	ret,thresh = cv2.threshold(filtered_mask,127,127,0)
	im2, contours, hierarchy = cv2.findContours(thresh,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)

	# take the convex hull of all of the contours
	convex_hulls = [cv2.convexHull(cnt) for cnt in contours]

	# find the biggest hull by area and extract it from the list of hulls
	maxArea = max([cv2.contourArea(cnt) for cnt in convex_hulls])
	contour = [cnt for cnt in convex_hulls if cv2.contourArea(cnt) >= maxArea][0]

	# find the minimum bounding rectangle
	rect = cv2.minAreaRect(contour)
	box = cv2.boxPoints(rect)
	box = np.int0(box)
	cv2.drawContours(img,[contour],-1,(0,255,0),3)

	# create a mask out of the bounding rectable
	contour_mask = np.zeros_like(img)
	cv2.drawContours(contour_mask, [box], 0, (255,255,255), cv2.FILLED)

	# Bitwise-AND mask and original image
	res = np.zeros_like(img) # Extract out the object and place into output image
	res[contour_mask == 255] = img[contour_mask == 255]

	############################################
	# Step 6
	# rotate the image to 90 degrees along the bottom
	############################################
	# compute the angle needed to square the rectangle along the bottom
	bottomLeft = (box[0][0], box[0][1])
	bottomRight = (box[3][0], box[3][1])

	# find the equation for the bottom line
	m = (1.0 * bottomLeft[1] - bottomRight[1]) / (1.0 * bottomLeft[0] - bottomRight[0])
	b = (bottomLeft[1]) - (m * bottomLeft[0])

	thirdPoint = (int((height - b) / m), height)
	cv2.line(res, thirdPoint, bottomRight, (255,255,255), 3)
	cv2.line(res, bottomRight, (bottomRight[0], height), (255,255,255), 3)
	cv2.line(res, thirdPoint, (bottomRight[0], height), (255,255,255), 3)

	hypLen = math.sqrt( math.pow( (thirdPoint[0]-bottomRight[0]), 2 ) + math.pow( (thirdPoint[1]-bottomRight[1]), 2 ) )
	oppLen = math.sqrt( 0 + math.pow( (bottomRight[1]-height), 2 ) ) 
	angle = math.sin(oppLen / hypLen)
	center = thirdPoint

	# rotate the image
	mat = cv2.getRotationMatrix2D(center, -math.degrees(angle), 1)
	rotated = cv2.warpAffine(res, mat, (width, height))
	rotated_mask = cv2.warpAffine(contour_mask, mat, (width, height))

	# write some more images
	if not video_mode:
		cv2.imwrite("contours.jpg", contour_mask)
		cv2.imwrite("contour_rotated.jpg", rotated_mask)
		cv2.imwrite("pre_rotation.jpg", res)
		cv2.imwrite("output.jpg", rotated)

	############################################
	# output the cropped rectangle to an image
	############################################


	return [hsv, mask, filtered_mask, im2, res, img, rotated_mask, rotated]

if video_mode:
	cap = create_capture(video_source)
	print(cap)
	while True:
		ret, img = cap.read()
		imgs = processImage(img, True)

		i = 0
		for img in imgs:
			cv2.imshow('STEP %d' % i, img)
			i = i + 1

		ch = 0xFF & cv2.waitKey(1)
		if ch == 27:
			break
	cv2.destroyAllWindows()
else:
	processImage(cv2.imread("test3.jpg"))