def classifyOne(self, depthImg, origImg, hsvImg, kp):
rows, cols, _ = origImg.shape
depthRect = img_manip.getBoundingRectangle(depthImg, kp, cols, rows, self.DEPTH_BOX_SZ)
material = self.classifyMaterial(depthRect)
if self.collect == 'material':
ratio = img_manip.nanRatioDepth(img)
self.dataStr = self.dataStr + ('%f %f' % ratio) + '\n'
boxHsv = img_manip.getBoundingRectangle(hsvImg, kp, cols, rows, self.COLOR_BOX_SZ)
if self.collect == 'color':
dataCol.append(boxHsv)
if material == 'wooden':
color = self.classifyWooden(boxHsv)
if color == 'green':
return GREEN_CUBE
if color == 'blue':
return BLUE_CUBE
houghingImg = img_manip.getBoundingRectangle(hsvImg, kp, cols, rows, self.HOUGH_BOX_SZ)
circleCount = img_manip.houghCountCircles(houghingImg, self.cannyParam, self.param1)
if circleCount >= 1:
return YELLOW_BALL if color == 'yellow' else RED_BALL
return YELLOW_CUBE if color == 'yellow' else RED_CUBE
else:
color = self.classifyPlastic(boxHsv)
if color == 'blue':
return BLUE_TRIANGLE
if color == 'green':
return GREEN_CYLINDER
if color == 'orange':
return PATRIC
if color == 'red':
return RED_HOLLOW_CUBE
boxCanny = img_manip.getBoundingRectangle(origImg, kp, cols, rows, self.CANNY_BOX_SZ)
if self.collect == 'purple':
self.dataColl.append(boxCanny)
shape = self.classifyPurple(boxCanny)
return PURPLE_STAR if shape == 'star' else PURPLE_CROSS
return '##### NOT WORKING #####' # Never gets here
def classifyMaterial(self, img): ratio = img_manip.nanRatioDepth(img) p = self.materialClf.predict(ratio)[0] predictedClass = self.materialClasses[p] return predictedClass