Files
opencv/beadDetect.py
T
2016-09-18 08:25:19 +00:00

278 lines
8.0 KiB
Python
Executable File

# import the necessary packages
from __future__ import print_function
from imutils.video import FPS
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)
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 = []
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 (perimeter > 0):
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")
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 = []
# 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()
gray = cv2.cvtColor(frame,cv2.COLOR_BGR2GRAY)
gray_blurred = cv2.medianBlur(gray,5)
# Canny edge detection
img1_canny = cv2.Canny(gray_blurred, 100, 50)
kernel = np.ones((1,1),np.uint8)
# img1_canny = cv2.blur(img1_canny,(3,3))
# img1_canny = cv2.dilate(img1_canny,kernel,iterations = 1)
# 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[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:3]})
# 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'])
# cv2.imshow('Thresholded',img1_thr)
cv2.imshow('Contours', img1_contours)
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()))
# 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]/255
green[i] = color[1]/255
red[i] = color[2]/255
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(Z,10,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()
# do a bit of cleanup
cv2.destroyAllWindows()