# TrainAndTest.py

import cv2 import numpy as np import operator import os

module level variables

MIN_CONTOUR_AREA = 100

RESIZED_IMAGE_WIDTH = 20 RESIZED_IMAGE_HEIGHT = 30

#

class ContourWithData():

# member variables ############################################################################
npaContour = None           # contour
boundingRect = None         # bounding rect for contour
intRectX = 0                # bounding rect top left corner x location
intRectY = 0                # bounding rect top left corner y location
intRectWidth = 0            # bounding rect width
intRectHeight = 0           # bounding rect height
fltArea = 0.0               # area of contour

def calculateRectTopLeftPointAndWidthAndHeight(self):               # calculate bounding rect info
    [intX, intY, intWidth, intHeight] = self.boundingRect
    self.intRectX = intX
    self.intRectY = intY
    self.intRectWidth = intWidth
    self.intRectHeight = intHeight

def checkIfContourIsValid(self):                            # this is oversimplified, for a production grade program
    if self.fltArea < MIN_CONTOUR_AREA: return False        # much better validity checking would be necessary
    return True

#

def main(): allContoursWithData = [] # declare empty lists, validContoursWithData = [] # we will fill these shortly

try:
    npaClassifications = np.loadtxt("classifications.txt", np.float32)                  # read in training classifications
except:
    print ("error, unable to open classifications.txt, exiting program\n")
    os.system("pause")
    return
# end try

try:
    npaFlattenedImages = np.loadtxt("flattened_images.txt", np.float32)                 # read in training images
except:
    print ("error, unable to open flattened_images.txt, exiting program\n")
    os.system("pause")
    return
# end try

npaClassifications = npaClassifications.reshape((npaClassifications.size, 1))       # reshape numpy array to 1d, necessary to pass to call to train

kNearest = cv2.ml.KNearest_create()                   # instantiate KNN object

kNearest.train(npaFlattenedImages, cv2.ml.ROW_SAMPLE, npaClassifications)

imgTestingNumbers = cv2.imread("test1.png")          # read in testing numbers image

if imgTestingNumbers is None:                           # if image was not read successfully
    print ("error: image not read from file \n\n")        # print error message to std out
    os.system("pause")                                  # pause so user can see error message
    return                                              # and exit function (which exits program)
# end if

imgGray = cv2.cvtColor(imgTestingNumbers, cv2.COLOR_BGR2GRAY)       # get grayscale image
imgBlurred = cv2.GaussianBlur(imgGray, (5,5), 0)                    # blur

                                                    # filter image from grayscale to black and white
imgThresh = cv2.adaptiveThreshold(imgBlurred,                           # input image
                                  255,                                  # make pixels that pass the threshold full white
                                  cv2.ADAPTIVE_THRESH_GAUSSIAN_C,       # use gaussian rather than mean, seems to give better results
                                  cv2.THRESH_BINARY_INV,                # invert so foreground will be white, background will be black
                                  11,                                   # size of a pixel neighborhood used to calculate threshold value
                                  2)                                    # constant subtracted from the mean or weighted mean

imgThreshCopy = imgThresh.copy()        # make a copy of the thresh image, this in necessary b/c findContours modifies the image

imgContours, npaContours, npaHierarchy = cv2.findContours(imgThreshCopy,             # input image, make sure to use a copy since the function will modify this image in the course of finding contours
                                             cv2.RETR_EXTERNAL,         # retrieve the outermost contours only
                                             cv2.CHAIN_APPROX_SIMPLE)   # compress horizontal, vertical, and diagonal segments and leave only their end points

for npaContour in npaContours:                             # for each contour
    contourWithData = ContourWithData()                                             # instantiate a contour with data object
    contourWithData.npaContour = npaContour                                         # assign contour to contour with data
    contourWithData.boundingRect = cv2.boundingRect(contourWithData.npaContour)     # get the bounding rect
    contourWithData.calculateRectTopLeftPointAndWidthAndHeight()                    # get bounding rect info
    contourWithData.fltArea = cv2.contourArea(contourWithData.npaContour)           # calculate the contour area
    allContoursWithData.append(contourWithData)                                     # add contour with data object to list of all contours with data
# end for

for contourWithData in allContoursWithData:                 # for all ...