opencv/beadDetect.py

# import the necessary packages
from __future__ import print_function

from imutils.video import FPS

import sys
import numpy as np
import argparse
import time
import cv2
import pprint
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

width = 320
height = 240

def colordistance(color1, color2):
	
	d0 = (color1[0]-color2[0])
	d1 = (color1[1]-color2[1])
	d2 = (color1[2]-color2[2])
	return np.math.sqrt(d0*d0 + d1*d1 + d2*d2)

def mindistance(colorList, color):
	minDist = 1000
	minIndex = 0
	index = 0
	for c in colorList:
		dist = colordistance(c, color)
		if dist < minDist:
			minDist = dist
			minIndex = index
			
		index += 1
		
	return (minDist, minIndex)
	
class FindObjects():
	def __init__(self):
		self.innerOuterRatio = 0.6	# const
		self.meanThickness = 0
		self.meanOuterDiameter = 0
		self.meanInnerDiameter = 0
		self.maxDiameter = -10000
		self.minDiameter = +10000
		self.pp = pprint.PrettyPrinter(indent=4)

	@staticmethod
	def getFamily(family, parentId, hierachy):
	
		family.append(parentId)
		if hierachy[0][parentId][2] != -1:
			return FindObjects.getFamily(family, hierachy[0][parentId][2], hierachy)
		else:
			return family


	def find(self, contours, hierachy):
		count = 0
		objects = []
		temp_objects = []
		objectsCandIds = []
		
		if hierachy is not None:
			while count != -1:
				family = FindObjects.getFamily([], count, hierachy)
	
				dupDict = {}
				members = []
				for member in family:
					x,y,w,h = cv2.boundingRect(contours[member]);
					key = str([x,y,w,h]) + '.key'				
					if not key in dupDict:
						dupDict[key] = member
						members.append({ 'id' : member, 'bbox' : [x,y,w,h]})
	
				objectsCandIds.append(members)
	
				count = hierachy[0][count][0]
		
			for family in objectsCandIds:
				obj = []
				for member in family:
					area = cv2.contourArea(contours[member['id']])
					perimeter = cv2.arcLength(contours[member['id']], True)
					pi = 3.14159265359
					Q = 0
					if (area > 10):
						Q = 4*pi*area/(perimeter*perimeter)
	
					if Q > 0.7:
						diameter = max(member['bbox'][2], member['bbox'][3])
						self.maxDiameter = max(self.maxDiameter, diameter)
						self.minDiameter = min(self.minDiameter, diameter)
						member['diameter'] = diameter
						member['pos'] = (int(member['bbox'][0] + member['bbox'][2]/2), int(member['bbox'][1] + member['bbox'][3]/2))
						obj.append(member)
	
				if obj:
					temp_objects.append(obj)
	
			for obj in temp_objects:
				for member in obj:
					if member['diameter'] < self.innerOuterRatio*self.maxDiameter:
						member['isHole'] = True
					else:
						member['isHole'] = False
	
				if len(obj) == 2:
					self.meanThickness = int(abs(obj[0]['diameter'] - obj[1]['diameter'])/2)
	
				objects.append({'thickness' : self.meanThickness, 'members' : obj})
						
		return objects

	def printStats(self):
		print ("maxDiameter   = " + str(self.maxDiameter) + " px")
		print ("minDiameter   = " + str(self.minDiameter) + " px")
		print ("meanThickness = " + str(self.meanThickness) + " px")
		
# construct the argument parse and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-n", "--num-frames", type=int, default=100,
	help="# of frames to loop over for FPS test")
ap.add_argument("-r", "--framerate", type=int, default=30,
	help="# of frames to loop over for FPS test")
ap.add_argument("-f", "--filename", type=str, default='piCamera',
	help="Whether or not frames should be displayed")
ap.add_argument("-c", "--calibrate", type=str, default='off',
	help="Whether calibration shall be performed")
ap.add_argument("-u", "--use-calibration", type=str, default='off',
	help="Whether calibration shall be used")
args = vars(ap.parse_args())

videoFile = args["filename"]
framerate = args["framerate"]

# created a *threaded *video stream, allow the camera sensor to warmup,
# and start the FPS counter
if videoFile == "piCamera":
	from PiVideoStream import PiVideoStream
	from picamera.array import PiRGBArray
	from picamera import PiCamera
	vs = PiVideoStream(resolution=(width,height), framerate=framerate).start()
	time.sleep(5.0)
else:
	vs = cv2.VideoCapture(videoFile)
	# skip first 5 seconds
	for i in range(0, 5*framerate):
		ret, frame = vs.read()

print("[INFO] sampling THREADED frames from `" + videoFile + "` at " + str(framerate) + " frame/s")
fps = FPS().start()
findObjects = FindObjects()

beadColors = []
numColorClasses = 0

if args["calibrate"] == 'on':
	lens_corr = np.ones((height, width, 3), np.float32)
	if videoFile == "piCamera":
		frame = vs.read()
	else:
		ret, frame = vs.read()
	
	Z = np.float32(frame)
	ref_color = cv2.mean(Z[int(height/2-8):int(height/2+8), int(width/2-8):int(width/2+8), :])
	print (ref_color)

	lens_corr = cv2.divide(ref_color, Z)

	params = list()
	params.append(cv2.IMWRITE_PNG_COMPRESSION)
	params.append(0)

	cv2.imwrite("lens_corr.png", np.uint8(lens_corr*127.0), params)

	sys.exit()


frame = np.ones((height, width, 3), np.float32)
lens_corr = np.float32(cv2.imread("lens_corr.png"))/140.0

# loop over some frames...this time using the threaded stream
while fps._numFrames < args["num_frames"]:
	# grab the frame from the threaded video stream

	if videoFile == "piCamera":
		frame = vs.read()
	else:
		ret, frame = vs.read()

	if args["use_calibration"] == 'on':
		frame = np.uint8(cv2.multiply(lens_corr, np.float32(frame)))

	kernel = np.ones((3,3),np.uint8)
	frame_dilated = cv2.dilate(frame,kernel,iterations = 1)
	gray = cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)
	gray_blurred = cv2.medianBlur(gray,5)

	# Canny edge detection
	img1_canny = cv2.Canny(gray_blurred, 100, 50)

	# Thresholding
#	ret,img1_thr = cv2.threshold(gray_blurred,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
#	img1_thr = cv2.adaptiveThreshold(gray_blurred,255,cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY,11,2)
#	img1_thr = cv2.adaptiveThreshold(gray_blurred,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,11,2)

	# Contours
	(_, contours, hierachy) = cv2.findContours(img1_canny.copy(),cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
	img1_contours = np.zeros((height,width,3), np.uint8)
	img1_contours = cv2.drawContours(img1_contours, contours, -1, (0,255,0), 1)
	
	objects = findObjects.find(contours, hierachy)
	
#	print (ids)

	img1_objects = frame.copy()
	img1_colors = np.zeros((height,width,3), np.uint8)
	maskCenter = 4
	roides = []
	beads = []
	for obj in objects:
		memberCount = 0
		radius = 0
		roi = 0
		for member in obj['members']:
			# Draw bounding box
			x,y,w,h = [member['bbox'][0], member['bbox'][1], member['bbox'][2], member['bbox'][3]];
			if member['isHole']:
				img1_objects = cv2.rectangle(img1_objects,(x,y),(x+w,y+h),(255,0,0),2)
				thickness = obj['thickness']-maskCenter
				diameter = member['diameter']+thickness+maskCenter
				roi = frame_dilated[y-thickness:y+diameter, x-thickness:x+diameter]
				pos = member['pos']
				if diameter > 0:
					mask = np.zeros((diameter,diameter,1), np.uint8)
					mask = cv2.circle(mask,(int(diameter/2), int(diameter/2)),radius,255,thickness)
					mean_color = cv2.mean(roi)
					img1_colors = cv2.circle(img1_colors,(int(pos[0]),int(pos[1])),int(diameter/2),mean_color,-1)
					beads.append({'pos' : pos, 'diameter' : diameter, 'color' : [mean_color[0]/255, mean_color[1]/255, mean_color[2]/255]})
#					img1_objects = cv2.circle(img1_objects,member['pos'],int(diameter/2),(255,255,255),thickness)
			else:
				img1_objects = cv2.rectangle(img1_objects,(x,y),(x+w,y+h),(0,0,255),2)

		memberCount += 1

	for bead in beads:
		beadColors.append(bead['color'])
			
			
	colorClasses = []
	for bead in beads:
		if not colorClasses:
			colorClasses.append(bead['color'])
		else:
			d, i = mindistance(colorClasses, bead['color'])
			if d > 0.14:
				colorClasses.append(bead['color'])
			else:
				colorClasses[i][0] = 0.5*colorClasses[i][0] + 0.5*bead['color'][0]
				colorClasses[i][1] = 0.5*colorClasses[i][1] + 0.5*bead['color'][1]
				colorClasses[i][2] = 0.5*colorClasses[i][2] + 0.5*bead['color'][2]
				
	if numColorClasses  != len(colorClasses):
		numColorClasses  = len(colorClasses)
		print("Found " + str(numColorClasses) + " color classes")
	
	
	cv2.imshow('Corrected', frame)
	cv2.imshow('Contours', frame_dilated)
	cv2.imshow('Colors', img1_colors)
	cv2.imshow('Canny',img1_canny)
	cv2.imshow('Objects',img1_objects)
	if videoFile == "piCamera":	
		cv2.waitKey(1)
	else:
#		cv2.waitKey(0)
		cv2.waitKey(int(1000.0/framerate))
		

	# update the FPS counter
	fps.update()

# stop the timer and display FPS information
fps.stop()
if videoFile == "piCamera":
	vs.stop()

print("[INFO] elasped time: {:.2f}".format(fps.elapsed()))
print("[INFO] approx. FPS: {:.2f}".format(fps.fps()))
if videoFile == "piCamera":
	print("[INFO] Camera : elasped time: {:.2f}".format(vs.getfps().elapsed()))
	print("[INFO] Camera: approx. FPS: {:.2f}".format(vs.getfps().fps()))

if numColorClasses:
	
	# Output stats
	findObjects.printStats()

	# Analyze colors
	numObjects = len(beadColors);
	blue = np.zeros(numObjects)
	green = np.zeros(numObjects)
	red = np.zeros(numObjects)
	plotColors = np.zeros((numObjects,3))
	i = 0
	for color in beadColors:
		blue[i] = color[0]
		green[i] = color[1]
		red[i] = color[2]
		plotColors[i] = [red[i], green[i], blue[i]]
		i += 1
	
	# Kmeans create color classes
	Z = np.float32(plotColors)
	
	# Define criteria = ( type, max_iter, epsilon )
	criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 100, 0.1)
	ret,label,center=cv2.kmeans(np.float32(plotColors), numColorClasses, None, criteria, 10, cv2.KMEANS_PP_CENTERS)
	
	print (center)
	
	cnt1 = 0
	for c1 in center:
		cnt2 = 0
		for c2 in center:
			print ("Color distance["+ str(cnt1) + "," + str(cnt2) + "] = " + str(colordistance(c1, c2)))
			cnt2 += 1
		cnt1 += 1	
		
	# Scatter plot of colors
	fig = plt.figure()
	ax = fig.add_subplot(111, projection='3d')
	ax.set_xlabel('blue')
	ax.set_ylabel('green')
	ax.set_zlabel('red')
	ax.scatter(blue, green, red, zdir='z', s=10, c=(0,0,0), lw = 0, depthshade=True)
	ax.scatter(center[:,2], center[:,1], center[:,0], zdir='z', s=500, facecolors=center, lw = 1, depthshade=True)
	
	plt.show()

else:
	print("Sorry, no colors have been detected")

# do a bit of cleanup
cv2.destroyAllWindows()