import sys
import numpy as np
import cv2

MIN_CONTOUR_AREA = 100

RESIZED_IMAGE_WIDTH = 20
RESIZED_IMAGE_HEIGHT = 30

imgTrainingNumbers = cv2.imread("training_chars.png")

imgGray = cv2.cvtColor(imgTrainingNumbers, cv2.COLOR_BGR2GRAY)
imgBlurred = cv2.GaussianBlur(imgGray, (5,5), 0)
imgThresh = cv2.adaptiveThreshold(imgBlurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2)

cv2.imshow("imgThresh", imgThresh)

imgThreshCopy = imgThresh.copy()
imgContours, _ = cv2.findContours(imgThreshCopy, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

npaFlattenedImages = np.empty((0, RESIZED_IMAGE_WIDTH * RESIZED_IMAGE_HEIGHT))
intClassifications = []

intValidChars = [ord('0'), ord('1'), ord('2'), ord('3'), ord('4'), ord('5'), ord('6'), ord('7'), ord('8'), ord('9'),
                 ord('A'), ord('B'), ord('C'), ord('D'), ord('E'), ord('F'), ord('G'), ord('H'), ord('I'), ord('J'),
                 ord('K'), ord('L'), ord('M'), ord('N'), ord('O'), ord('P'), ord('Q'), ord('R'), ord('S'), ord('T'),
                 ord('U'), ord('V'), ord('W'), ord('X'), ord('Y'), ord('Z')]

for npaContour in imgContours:
    if cv2.contourArea(npaContour) > MIN_CONTOUR_AREA:
        [intX, intY, intW, intH] = cv2.boundingRect(npaContour)

        cv2.rectangle(imgTrainingNumbers, (intX, intY), (intX+intW, intY+intH), (0, 0, 255), 2)

        imgROI = imgThresh[intY:intY+intH, intX:intX+intW]
        imgROIResized = cv2.resize(imgROI, (RESIZED_IMAGE_WIDTH, RESIZED_IMAGE_HEIGHT))

        cv2.imshow("imgROI", imgROI)
        cv2.imshow("imgROIResized", imgROIResized)
        cv2.imshow("training_numbers.png", imgTrainingNumbers)

        intChar = cv2.waitKey(0)
        if intChar >= 97:
            intChar -= 32

        if intChar == 27:
            sys.exit()
        elif intChar in intValidChars:
            intClassifications.append(intChar)
            npaFlattenedImage = imgROIResized.reshape((1, RESIZED_IMAGE_WIDTH * RESIZED_IMAGE_HEIGHT))  # flatten image to 1d numpy array so we can write to file later
            npaFlattenedImages = np.append(npaFlattenedImages, npaFlattenedImage, 0)                    # add current flattened impage numpy array to list of flattened image numpy arrays

fltClassifications = np.array(intClassifications, np.float32)                   # convert classifications list of ints to numpy array of floats
npaClassifications = fltClassifications.reshape((fltClassifications.size, 1))   # flatten numpy array of floats to 1d so we can write to file later

print("\ntraining complete !!\n")

np.savetxt("classifications.txt", npaClassifications)
np.savetxt("flattened_images.txt", npaFlattenedImages)

cv2.destroyAllWindows()



