formatImages.py

'''
Created on 4 Jul 2017

@author: pings
'''
# for file manipulation
import os, numpy as np, collections
from PIL import Image, ImageChops, ImageFilter

Datasets = collections.namedtuple('Datasets', ['train', 'test'])

# number of images of each letter saved
DATAPOINTS = 10000 # no. generated by gen_data()
LETTERS = 24 # no of letters outputted

class Dataset:
    def __init__(self, images, labels, letters = LETTERS):
        self._size = len(images) # or len(labels)
        self._images = images
        self._labels = labels
        self._letters = letters
        # starting point for the get_batch function
        self._batchIndex = 0
        
    def next_batch(self, batch_size, shuffle = True):
        # return a numpy tuple of 2 arrays, shape (batch_size, 45*45) and (batch_size, _ketters)
        
        batchImages = np.empty([batch_size, 45*45])
        batchLabels = np.empty([batch_size, self._letters])
        
        if (shuffle):
            import random
            # pick a random batch
            for i in range(batch_size):
                j = random.randint(0, self._size-1)
                batchImages[i, :] = self._images[j, :]
                batchLabels[i, :] = self._labels[j, :]
        else:
            start = self._batchIndex
            
            if (self._batchIndex + batch_size) < self._size: # return a batch as normal
                batchImages = self._images[start:(start+batch_size), :]
                batchLabels = self._labels[start:(start+batch_size), :]
                self._batchIndex = start + batch_size
            else: # wraps around back to the beginning
                batchImages = np.append(self._images[start:self._size, :],
                                        self._images[0:(start + batch_size)%self._size, :],
                                        axis = 0)
                batchLabels = np.append(self._labels[start:self._size, :],
                                        self._labels[0:(start + batch_size)%self._size, :],
                                        axis = 0)
                self._batchIndex = (start + batch_size) % self._size
        
        return (batchImages, batchLabels)
    
    @property
    def images(self):
        return self._images
    
    @property
    def labels(self):
        return self._labels
                
def get_all_data(srcpath, points = DATAPOINTS):
    import re
    
    # reformats srcpath
    srcpath = srcpath.replace("\\", "/")
    
    ''' HUMAN SORTING '''
    def atoi(text):
        return int(text) if text.isdigit() else text

    def natural_keys(text):
        '''
        alist.sort(key=natural_keys) sorts in human order
        http://nedbatchelder.com/blog/200712/human_sorting.html
        (See Toothy's implementation in the comments)
        '''
        return [ atoi(c) for c in re.split('(\d+)', text) ]
    
    print(next(os.walk(srcpath))[1])
    # walk through the directories, each letter's data is stored in a folder
    letterDirs = next(os.walk(srcpath))[1]
    
    # sort according to human intuition, i.e. letters then numbers
    letterDirs.sort(key = natural_keys)
    
    # data will be saved here
    images = np.empty([points*LETTERS, 45*45], int) # points*18 images
    labels = np.empty([points*LETTERS], int) # correct answers of each image, convert to one-hot after
    
    # j represents the indices where fill our arrays
    j = 0
    
    # i represents the labels of the letters
    for i in range(LETTERS):
        for _, _, files in os.walk(srcpath + "/" + letterDirs[i]):
            # path to directory
            dirPath = srcpath + "/" + letterDirs[i] + "/"
            
            # c keeps count of the no. of letters taken
            c = 1
            for f in files:
                imgPath = dirPath + f
                
                # convert to greyscale
                img = Image.open(imgPath).convert('LA')
                imgArr = np.array(img)
                
                # greyscale - only interested in the first argument for each pixel
                # which is the greyness of the pixel
                # flatten converts to 1d array; squeeze turns into proper 2d array from 3d
                imgArr = np.squeeze(imgArr[:,:,0]).flatten()
                print(imgArr.shape)
                images[j, :] = 255 - imgArr
                labels[j] = i
                
                print("progress: " + str(j))
                j += 1
                
                # update count
                c += 1
                
                # loop will also naturally break if the files run out
                if (c > points) or (c > DATAPOINTS):
                    break;
                
    labels = np.eye(LETTERS)[labels] # converting to one-hot
    
    return images, labels

def get_data(srcpath, shuffle = True, points = DATAPOINTS):
    images, labels = get_all_data(srcpath, points = points)
    
    # appends the 2 sets of data together. data[:, :-LETTERS] is the images, data[:, -LETTERS:] is the labels (for shuffling)
    data = np.append(images, labels, axis = 1)
    if shuffle:
        np.random.shuffle(data)
        images = data[:, :-LETTERS]
        labels = data[:, -LETTERS:]
    
    # train with 90% of the data, and test with 10%
    cutoff = int(round(len(images)*0.9))
    train = Dataset(images[:cutoff, :], labels[:cutoff, :])
    test = Dataset(images[cutoff:, :], labels[cutoff:, :])
    
    return Datasets(train = train, test = test)

# code below processes data and then save in cleanData directory
def gen_data():
    # get all the .bmp files in a dir
    # images = glob.glob(os.path.dirname(__file__)+"\LETT_CAP_NORM.ALPHA\*.bmp")
    
    root = os.path.dirname(__file__)
    root_raw = os.path.dirname(__file__) + "/goodData"
    root = root.replace("\\", "/")
    root_raw = root_raw.replace("\\", "/")
    print(root, root_raw)
    
    # directory of letter folders 
    # format: the letters dir corresponds to the indices of the list (lexicographic)
    letterDirs = next(os.walk(root_raw))[1]
    print(letterDirs)
    
    # i corresponds to the one-hot array's index of the letters
    for i in range(len(letterDirs)):
        for _, _, files in os.walk(root_raw + "/" + letterDirs[i]):
            # every file is walked here of letter letterDirs[i], all .bmp
            
            # make the dir if one does not exist
            dirPath = root + "/" + "cleanData_new/" + str(i) + "/"
            if not os.path.exists(dirPath):
                os.makedirs(dirPath)
            
            # j keeps track of the filenames
            j = 1
            
            for f in files:
                if (f.endswith('.bmp')):
                    savePath = dirPath + str(j) + ".bmp"
                    
                    if not os.path.isfile(savePath):
                        # do the image processing here
                        imgpath = root_raw + "/" + letterDirs[i] + "/" + f
                        # im = Image.open(imgpath)
                        
                        # save im to file
                        # im.save(savePath)
                        import shutil
                        shutil.copy2(imgpath, savePath) # simply copy the image over. no processing.
                    # report to console
                    print("processed: " + str(j) + " letter = " + str(i))
                    
                    j += 1
                    
                    # we'll stop at DATAPOINTS images
                    if (j > DATAPOINTS):
                        break
                

# remove white space, and then fill on all sides to 45x45
def processImg(path):    
    # cleans up file path, remove later (already done in get_data)
    path = path.replace("\\", "/")
    
    # rids of all white spaces
    def trim(img):
        bg = Image.new(img.mode, img.size, img.getpixel((0,0)))
        diff = ImageChops.difference(img, bg)
        diff = ImageChops.add(diff, diff, 2.0, -100)
        bbox = diff.getbbox()
        if bbox:
            return img.crop(bbox)
    
    # reformats to a nice 45*45 with the character centred
    def addWhiteSpace(img):
        # firstly resize image to fill as much space as possible
        max_w, max_h = 45, 45
        img_w, img_h = img.size
        
        # calculates the ratio to resize by
        r = min(max_w / img_w, max_h / img_h)
        
        # then resizes the image
        newDims = (int(round(img_w*r)), int(round(img_h*r)))
        img = img.resize(newDims)
        img_w, img_h = newDims
        
        # create a background for our image to lay on top of
        bg = Image.new('RGBA', (45, 45), (255,255,255))
        bg_w, bg_h = bg.size
        
        # offset centres our image onto the white background
        offset = ((bg_w - img_w) // 2, (bg_h - img_h) // 2)
        
        bg.paste(img, offset)
        
        # add antialias
        bg = bg.filter(ImageFilter.GaussianBlur(radius = 1))
        return bg
    
    # we want to return this image
    image = Image.open(path)
    image = addWhiteSpace(trim(image))
    return image

# takes in array of an image
def processImage(image):    
    # rids of all white spaces
    def trim(img):
        bg = Image.new(img.mode, img.size, img.getpixel((0,0)))
        diff = ImageChops.difference(img, bg)
        diff = ImageChops.add(diff, diff, 2.0, -100)
        bbox = diff.getbbox()
        if bbox:
            return img.crop(bbox)
    
    # reformats to a nice 45*45 with the character centred
    def addWhiteSpace(img):
        # firstly resize image to fill as much space as possible
        max_w, max_h = 45, 45
        img_w, img_h = img.size
        
        # calculates the ratio to resize by
        r = min(max_w / img_w, max_h / img_h)
        
        # then resizes the image
        newDims = (int(round(img_w*r)), int(round(img_h*r)))
        img = img.resize(newDims)
        img_w, img_h = newDims
        
        # create a background for our image to lay on top of
        bg = Image.new('RGBA', (45, 45), (255,255,255))
        bg_w, bg_h = bg.size
        
        # offset centres our image onto the white background
        offset = ((bg_w - img_w) // 2, (bg_h - img_h) // 2)
        
        bg.paste(img, offset)
        
        # add antialias
        bg = bg.filter(ImageFilter.GaussianBlur(radius = 1))
        return bg
    
    # we want to return this image
    return addWhiteSpace(trim(image))

def toArray(image):
    # convert to greyscale
    imgArr = np.array(image)
    
    # greyscale - only interested in the first argument for each pixel
    # which is the greyness of the pixel
    # flatten converts to 1d array; squeeze turns into proper 2d array from 3d
    imgArr = np.reshape(imgArr[:,:,0], (1,45*45))
    print(imgArr.shape)
    return 255 - imgArr

def arrToImage(arr):
    # takes in a 2d array and converts it to image
    print("drawing...")
    from matplotlib import pyplot as plt
    plt.imshow(arr, cmap='gray', interpolation='nearest', vmin=0, vmax=255)
    plt.show()

if __name__ == '__main__':
    '''stuff = get_data("C:\\Users\\pings\\Desktop\\Python\\Tensorflow\\GreekLetters\\cleanData_new", points = 1000, shuffle = True)
    
    xs, ys = stuff.train.next_batch(100)
    
    i = 5
    testImage = xs[i].reshape((45,45))
    for j in range(LETTERS):
        if ys[i,j] == 1:
            print("label: " + str(j))
            
    arrToImage(testImage)'''