batch.py

def _load_image(filename) :
    cv2_arr = cv2.imread(filename)
    cv2_arr = rescale(cv2_arr,size=(224, 224))
    arr = np.asarray(cv2_arr, dtype="float32")
    return arr

def rescale(cv2_arr, size):
    return cv2.resize(cv2_arr, dsize=size, interpolation=cv2.INTER_CUBIC)

def load_images(img_paths):
    batch_images = []
    jpg_paths = list(filter(lambda x: ".jpg" in x, img_paths))
    for path in jpg_paths:
        img_arr = _load_image(path)
        batch_images.append(img_arr)
    return batch_images
        
def one_hot(n, num_classes):
    label = [0] * num_classes
    label[n] = 1
    return label

def split_image(images):
    batch, _, _, _ = images.shape
    labels = one_hot(8, 10)
    for i in range(batch - 1):
        labels = np.append(labels, one_hot(8, 10), axis=0)
    return (images, labels), (images, labels)


def stack_n(label, n):
    labels = []
    for i in range(n):
        labels.append(label)
    return labels


class Batch(object):
    idx = 0 
    remain = 0
    train_data = []
    train_label = []
    
    def __init__(self, train_dir="/home/soma03/projects/ys/codes/data/ir_ph1_v2/train/train_data"):
        self.train_dir = train_dir
        self.train_data, self.train_label = make_image_list(self.train_dir)
        self.remain = len(self.train_data)
    
    @staticmethod
    def make_image_list(train_dir):
        path = os.listdir(train_dir)
        path = list(map(lambda x: os.path.join(train_dir, x), path))
        train_data = []
        train_label = []

        idx = 0
        for folder in path: 
        #     print("directory [" + os.path.basename(folder) + "]")
        #     print(os.listdir(folder))
        #     print()
            files = os.listdir(folder)
            files_path = list(map(lambda file : os.path.join(folder, file), files))
            class_data = load_images(files_path)
            num_img = len(class_data)
            #nsml에서는 이코드를 사용
        #     label = one_hot(int(os.path.basename(folder)), num_classes)    
            label = one_hot(idx, num_classes)    
            class_label = stack_n(label, num_img)
            train_data += class_data
            train_label += class_label
            idx += 1
        return (train_data, train_label)
    
    def next_batch(self, n):
        combined = list(zip(self.train_data, self.train_label))
        shuffle(combined)
        data, label = zip(*combined)
        return data[:n], label[:n]

#     def next_batch(self, n):
#         if n <= self.remain:
#             self.remain -= n
#             self.idx += n
#             return self.train_data[self.idx: self.idx + n], self.train_label[self.idx: self.idx + n]
#         elif n > self.remain:
#             self.idx += n
#             self.remain = 0
#             return self.train_data[self.idx : self.remain], self.train_label[self.idx: self.remain]