forked from danielvarga/keras-finetuning
-
Notifications
You must be signed in to change notification settings - Fork 0
/
new_batch_gen.py
43 lines (39 loc) · 1.21 KB
/
new_batch_gen.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
import numpy as np
import sys, glob
import cv2
import os
import argparse
import hashlib
import os.path
import scipy
import scipy.misc
import random
from keras.utils import np_utils
def getImageClassifierFromDirGen(imagesDir,width=224,height=224,batchSize=32,augment=False):
allImages = glob.glob(os.path.join(imagesDir,"*/*.jpg")) + glob.glob( os.path.join( imagesDir,"*/*.jpeg")) + glob.glob(os.path.join(imagesDir,"*/*.png"))
random.shuffle(allImages)
classes = set()
for image in allImages :
className = image.split('/')[-2]
classes.add(className)
classes = list(classes)
nClasses = len(classes)
classesIds = dict( enumerate(classes))
classesIds = {v: k for k, v in classesIds.iteritems()}
X_batch = []
Y_batch = []
while True :
for image in allImages :
X_batch.append( getImageVec(image, width=width, height=height, augment=augment))
className = image.split('/')[-2]
classId = classesIds[className]
classVec = np.zeros(( nClasses ) )
classVec[classId] = 1
Y_batch.append( classVec )
if len( X_batch ) == batchSize:
tx , ty = X_batch , Y_batch
X_batch = []
Y_batch = []
yield np.array(tx), np.array(ty)
if __name__ == "__main__":
getImageClassifierFromDirGen(sys.argv[1])