#pip install tqdm

import os
import cv2
import numpy as np
from tqdm import tqdm

REBUILD_DATA = False

class Currency():
    IMG_SIZE = 150
    BACK = "C:/Users/Aayush/Desktop/Fake-Currency-Detection/Images/back"
    FRONT = "C:/Users/Aayush/Desktop/Fake-Currency-Detection/Images/front"
    LABELS = {BACK: 0, FRONT: 1}
    training_data = []
    backcount = 0
    frontcount = 0
    
    def make_training_data(self):
        for label in self.LABELS:
            print(label)
            for f in tqdm(os.listdir(label)):
                try:
                    path = os.path.join(label,f)
                    img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
                    img = cv2.resize(img, (self.IMG_SIZE, self.IMG_SIZE))
                    self.training_data.append([np.array(img), np.eye(2)[self.LABELS[label]]])


                    if label == self.FRONT:
                        self.frontcount += 1
                    elif label == self.BACK:
                        self.backcount += 1
                except Exception as e:
                    pass
                    #print(str(e))
        
        
            
        np.random.shuffle(self.training_data)
        np.save("training_data.npy",self.training_data)
        print('Front:', self.frontcount)
        print('Back:', self.backcount)
    
if  REBUILD_DATA:
    currency = Currency()
    currency.make_training_data()
                           

training_data = np.load("training_data.npy", allow_pickle=True)

print(len(training_data))