def __init__(self, params, indexlist, phase, proc_id):
super(BatchLoader, self).__init__()
self.indexlist = indexlist
self.proc_id = proc_id
self.batch_size = params['batch_size']
self.im_shape = params['im_shape']
self.phase = phase
self.queue = Queue(_QSIZE)
#rec_conn, send_conn = Pipe()
# self.rec_conn = rec_conn
# self.send_conn = send_conn
## Dividing with rest the batch size for the jobs we have
self.batch_ck_size = self.batch_size // _nJobs
## in case of the last jobs adding the rest
if self.proc_id == (_nJobs - 1):
self.batch_ck_size += self.batch_size % _nJobs
## Opening LMDB
lmdb_output_pose_env = lmdb.Environment(params['source'] +
'/pose_lmdb/',
readonly=True,
lock=False)
self.cur_pose = lmdb_output_pose_env.begin().cursor()
lmdb_output_flip_env = lmdb.Environment(params['source'] +
'/flip_lmdb/',
readonly=True,
lock=False)
self.cur_flip = lmdb_output_flip_env.begin().cursor()
lmdb_output_land_env = lmdb.Environment(params['source'] +
'/land_lmdb/',
readonly=True,
lock=False)
self.cur_land = lmdb_output_land_env.begin().cursor()
################
self.Nimgs = len(self.indexlist)
# this class does some simple data-manipulations
#proto_data = open(params['mean_file'], "rb").read()
#a = caffe.io.caffe_pb2.BlobProto.FromString(proto_data)
#mean = caffe.io.blobproto_to_array(a)[0]
## mean is read BGR and c,h,w; we convert it to h,w,c.
## BGR is OK since OpenCV and caffe are BGR
## Then MySimpleTransformer will remove mean after that the image
## has been changed to BGR as well. So apple-to-apple.
self.transformer = MySimpleTransformer()
self.aug_tr = aug_tracker.AugmentationTracker()
if params['mean_file'] is not None:
mean = np.load(params['mean_file'])
mean = mean.transpose(1, 2, 0)
mean = np.float32(mean)
self.transformer.set_mean(mean)
if self.phase == 1:
util.myprint("BatchLoader_valid" + str(self.proc_id) +
" initialized with " + str(self.Nimgs) + " images")
else:
util.myprint("BatchLoader_train" + str(self.proc_id) +
" initialized with " + str(self.Nimgs) + " images")
util.myprint("This will process: " + str(self.batch_ck_size) +
'/' + str(self.batch_size))
class BatchLoader(Process):
"""
This class abstracts away the loading of images.
Images can either be loaded singly, or in a batch. The latter is used for
the asyncronous data layer to preload batches while other processing is
performed.
"""
def __init__(self, params, indexlist, phase, proc_id):
super(BatchLoader, self).__init__()
self.indexlist = indexlist
self.proc_id = proc_id
self.batch_size = params['batch_size']
self.im_shape = params['im_shape']
self.phase = phase
self.queue = Queue(_QSIZE)
#rec_conn, send_conn = Pipe()
# self.rec_conn = rec_conn
# self.send_conn = send_conn
## Dividing with rest the batch size for the jobs we have
self.batch_ck_size = self.batch_size//_nJobs
## in case of the last jobs adding the rest
if self.proc_id == (_nJobs - 1):
self.batch_ck_size += self.batch_size % _nJobs
## Opening LMDB
lmdb_output_pose_env = lmdb.Environment(params['source']+'/pose_lmdb/', readonly=True, lock=False)
self.cur_pose = lmdb_output_pose_env.begin().cursor()
lmdb_output_flip_env = lmdb.Environment(params['source']+'/flip_lmdb/', readonly=True, lock=False)
self.cur_flip = lmdb_output_flip_env.begin().cursor()
lmdb_output_land_env = lmdb.Environment(params['source']+'/land_lmdb/', readonly=True, lock=False)
self.cur_land = lmdb_output_land_env.begin().cursor()
################
self.Nimgs = len(self.indexlist)
# this class does some simple data-manipulations
#proto_data = open(params['mean_file'], "rb").read()
#a = caffe.io.caffe_pb2.BlobProto.FromString(proto_data)
#mean = caffe.io.blobproto_to_array(a)[0]
## mean is read BGR and c,h,w; we convert it to h,w,c.
## BGR is OK since OpenCV and caffe are BGR
## Then MySimpleTransformer will remove mean after that the image
## has been changed to BGR as well. So apple-to-apple.
self.transformer = MySimpleTransformer()
self.aug_tr = aug_tracker.AugmentationTracker()
if params['mean_file'] is not None:
mean = np.load(params['mean_file'])
mean = mean.transpose(1, 2, 0)
mean = np.float32(mean)
self.transformer.set_mean(mean)
if self.phase == 1:
util.myprint("BatchLoader_valid" + str(self.proc_id) + " initialized with " + str(self.Nimgs) +" images")
else:
util.myprint("BatchLoader_train" + str(self.proc_id) + " initialized with " + str(self.Nimgs) +" images")
util.myprint("This will process: " + str(self.batch_ck_size)+'/'+str(self.batch_size) )
def run(self):
if self.phase == 1:
util.myprint("Process started pre-fetching for Validation " + str(self.proc_id) + " : nimgs " + str(self.Nimgs) )
else:
util.myprint("Process started pre-fetching for Training " + str(self.proc_id) + " : nimgs " + str(self.Nimgs) )
## Counter to the entire augmented set
count = 0
## Counter to the relative mini-batch
countStep = 0
## Pre-allocate the data for the mini-batch
listData = [None]*self.batch_ck_size
while True:
for ii in range(0,self.Nimgs):
####### Checking if we finished an (augmented) epoch
if count == self.Nimgs:
util.myprint("Finished an (augmented) epoch for loader id " + str(self.proc_id) + "...shuffling")
count = 0
shuffle(self.indexlist)
# ######## Part to resume/wait a certain process when the other is operating
# if self.phase == 1:
# if not _eventValidList[self.proc_id].is_set():
# util.myprint('Waiting Validation Loader ' + str(self.proc_id) + ' to start again')
# _eventValidList[self.proc_id].wait()
# else:
# if not _eventTrainList[self.proc_id].is_set():
# util.myprint('Waiting Train Loader ' + str(self.proc_id) + ' to start again')
# _eventTrainList[self.proc_id].wait()
### Starting to do augmentation
batch_img = None
#index is of form:
#blur_fr_13 XXXm.0hhvfrvXXX_MS000024 !!TMPDIR!!/imgs/XXXm.0hhvfrvXXX/XXXm.0hhvfrvXXX_MS000024.jpg 0
index = self.indexlist[ii]
index = index.split(' ')
aug_type = index[0] #augemntation type
image_key = index[1] # image key
image_file_name = index[2] #image
label = np.float32(index[3]) #label
## Loading the image with OpenCV
flipON = int( np.frombuffer( self.cur_flip.get(image_key) )[1] ) == 1
im = cv2.imread(image_file_name,cv2.CV_LOAD_IMAGE_COLOR)
## Check immediately if we have to flip an image
if flipON:
im = cv2.flip(im, 1)
im_arr = np.asarray(im)
aug_im = None
if 'align2d' in aug_type or 'blur' in aug_type:
lmark = self.cur_land.get(image_key)
lmark = np.frombuffer(lmark, dtype='float64').reshape(68,2)
lmarks = np.zeros((1,68,2))
lmarks[0] = lmark
aug_im = self.aug_tr.augment_fast(aug_type=aug_type,img=im,landmarks=lmarks,flipON=flipON)
elif 'render' in aug_type:
prj_matrix = np.frombuffer(self.cur_pose.get(image_key+'_'+aug_type), dtype='float64').reshape(3,4)
prj_matrix = np.asmatrix(prj_matrix)
aug_im = self.aug_tr.augment_fast(aug_type=aug_type,img=im,prj_matrix=prj_matrix,flipON=flipON)
try:
aug_im = cv2.resize(aug_im, ( self.im_shape[0], self.im_shape[1] ),\
interpolation=cv2.INTER_LINEAR )
batch_img = self.transformer.preprocess(aug_im)
except Exception as ex:
util.myprint("Warning: Was not able to use aug_img because: " + str(ex))
util.myprint( "Skipping the image: " + image_file_name)
count += 1
##If image have been processes correctly, add it to the mini-batch
if batch_img is not None:
data = {'img': batch_img , 'label' : label}
listData[countStep] = data
countStep+=1
if countStep == self.batch_ck_size:
isDone = False
while not isDone:
try:
##This mini-batch is ready to be sent for train
## Resetting the relative listData and countStep
self.queue.put_nowait( list(listData) )
except std_Queue.Full as full:
pass
else:
#self.send_conn.send( (listData) )
countStep = 0
isDone = True
listData = [None]*self.batch_ck_size