|
- import numpy as np
- import cv2
-
-
-
- class TrainImageReader:
- def __init__(self, imdb, im_size, batch_size=128, shuffle=False):
-
- self.imdb = imdb
- self.batch_size = batch_size
- self.im_size = im_size
- self.shuffle = shuffle
-
- self.cur = 0
- self.size = len(imdb)
- self.index = np.arange(self.size)
- self.num_classes = 2
-
- self.batch = None
- self.data = None
- self.label = None
-
- self.label_names= ['label', 'bbox_target', 'landmark_target']
- self.reset()
- self.get_batch()
-
- def reset(self):
- self.cur = 0
- if self.shuffle:
- np.random.shuffle(self.index)
-
- def iter_next(self):
- return self.cur + self.batch_size <= self.size
-
- def __iter__(self):
- return self
-
- def __next__(self):
- return self.next()
-
- def next(self):
- if self.iter_next():
- self.get_batch()
- self.cur += self.batch_size
- return self.data,self.label
- else:
- raise StopIteration
-
- def getindex(self):
- return self.cur / self.batch_size
-
- def getpad(self):
- if self.cur + self.batch_size > self.size:
- return self.cur + self.batch_size - self.size
- else:
- return 0
-
- def get_batch(self):
- cur_from = self.cur
- cur_to = min(cur_from + self.batch_size, self.size)
- imdb = [self.imdb[self.index[i]] for i in range(cur_from, cur_to)]
- data, label = get_minibatch(imdb)
- self.data = data['data']
- self.label = [label[name] for name in self.label_names]
-
-
-
- class TestImageLoader:
- def __init__(self, imdb, batch_size=1, shuffle=False):
- self.imdb = imdb
- self.batch_size = batch_size
- self.shuffle = shuffle
- self.size = len(imdb)
- self.index = np.arange(self.size)
-
- self.cur = 0
- self.data = None
- self.label = None
-
- self.reset()
- self.get_batch()
-
- def reset(self):
- self.cur = 0
- if self.shuffle:
- np.random.shuffle(self.index)
-
- def iter_next(self):
- return self.cur + self.batch_size <= self.size
-
- def __iter__(self):
- return self
-
- def __next__(self):
- return self.next()
-
- def next(self):
- if self.iter_next():
- self.get_batch()
- self.cur += self.batch_size
- return self.data
- else:
- raise StopIteration
-
- def getindex(self):
- return self.cur / self.batch_size
-
- def getpad(self):
- if self.cur + self.batch_size > self.size:
- return self.cur + self.batch_size - self.size
- else:
- return 0
-
- def get_batch(self):
- cur_from = self.cur
- cur_to = min(cur_from + self.batch_size, self.size)
- imdb = [self.imdb[self.index[i]] for i in range(cur_from, cur_to)]
- data= get_testbatch(imdb)
- self.data=data['data']
-
-
-
-
- def get_minibatch(imdb):
-
- # im_size: 12, 24 or 48
- num_images = len(imdb)
- processed_ims = list()
- cls_label = list()
- bbox_reg_target = list()
- landmark_reg_target = list()
-
- for i in range(num_images):
- im = cv2.imread(imdb[i]['image'])
- #im = Image.open(imdb[i]['image'])
-
- if imdb[i]['flipped']:
- im = im[:, ::-1, :]
- #im = im.transpose(Image.FLIP_LEFT_RIGHT)
-
- cls = imdb[i]['label']
- bbox_target = imdb[i]['bbox_target']
- landmark = imdb[i]['landmark_target']
-
- processed_ims.append(im)
- cls_label.append(cls)
- bbox_reg_target.append(bbox_target)
- landmark_reg_target.append(landmark)
-
- im_array = np.asarray(processed_ims)
-
- label_array = np.array(cls_label)
-
- bbox_target_array = np.vstack(bbox_reg_target)
-
- landmark_target_array = np.vstack(landmark_reg_target)
-
- data = {'data': im_array}
- label = {'label': label_array,
- 'bbox_target': bbox_target_array,
- 'landmark_target': landmark_target_array
- }
-
- return data, label
-
-
- def get_testbatch(imdb):
- assert len(imdb) == 1, "Single batch only"
- im = cv2.imread(imdb[0]['image'])
- data = {'data': im}
- return data
|