Spaces:
Runtime error
Runtime error
| import torchvision.transforms as transforms | |
| import numpy as np | |
| import os | |
| import cv2 | |
| def convert_image_to_tensor(image): | |
| """convert an image to pytorch tensor | |
| Parameters: | |
| ---------- | |
| image: numpy array , h * w * c | |
| Returns: | |
| ------- | |
| image_tensor: pytorch.FloatTensor, c * h * w | |
| """ | |
| transform = transforms.ToTensor() | |
| return transform(image) | |
| def convert_chwTensor_to_hwcNumpy(tensor): | |
| """convert a group images pytorch tensor(count * c * h * w) to numpy array images(count * h * w * c) | |
| Parameters: | |
| ---------- | |
| tensor: numpy array , count * c * h * w | |
| Returns: | |
| ------- | |
| numpy array images: count * h * w * c | |
| """ | |
| return np.transpose(tensor.detach().numpy(), (0,2,3,1)) | |
| class ImageDB(object): | |
| def __init__(self, image_annotation_file, prefix_path='', mode='train'): | |
| self.prefix_path = prefix_path | |
| self.image_annotation_file = image_annotation_file | |
| self.classes = ['__background__', 'face'] | |
| self.num_classes = 2 | |
| self.image_set_index = self.load_image_set_index() | |
| self.num_images = len(self.image_set_index) | |
| self.mode = mode | |
| def load_image_set_index(self): | |
| """Get image index | |
| Parameters: | |
| ---------- | |
| Returns: | |
| ------- | |
| image_set_index: str | |
| relative path of image | |
| """ | |
| assert os.path.exists(self.image_annotation_file), 'Path does not exist: {}'.format(self.image_annotation_file) | |
| with open(self.image_annotation_file, 'r') as f: | |
| image_set_index = [x.strip().split(' ')[0] for x in f.readlines()] | |
| return image_set_index | |
| def load_imdb(self): | |
| """Get and save ground truth image database | |
| Parameters: | |
| ---------- | |
| Returns: | |
| ------- | |
| gt_imdb: dict | |
| image database with annotations | |
| """ | |
| gt_imdb = self.load_annotations() | |
| return gt_imdb | |
| def real_image_path(self, index): | |
| """Given image index, return full path | |
| Parameters: | |
| ---------- | |
| index: str | |
| relative path of image | |
| Returns: | |
| ------- | |
| image_file: str | |
| full path of image | |
| """ | |
| index = index.replace("\\", "/") | |
| if not os.path.exists(index): | |
| image_file = os.path.join(self.prefix_path, index) | |
| else: | |
| image_file=index | |
| if not image_file.endswith('.jpg'): | |
| image_file = image_file + '.jpg' | |
| assert os.path.exists(image_file), 'Path does not exist: {}'.format(image_file) | |
| return image_file | |
| def load_annotations(self,annotion_type=1): | |
| """Load annotations | |
| Parameters: | |
| ---------- | |
| annotion_type: int | |
| 0:dsadsa | |
| 1:dsadsa | |
| Returns: | |
| ------- | |
| imdb: dict | |
| image database with annotations | |
| """ | |
| assert os.path.exists(self.image_annotation_file), 'annotations not found at {}'.format(self.image_annotation_file) | |
| with open(self.image_annotation_file, 'r') as f: | |
| annotations = f.readlines() | |
| imdb = [] | |
| for i in range(self.num_images): | |
| annotation = annotations[i].strip().split(' ') | |
| index = annotation[0] | |
| im_path = self.real_image_path(index) | |
| imdb_ = dict() | |
| imdb_['image'] = im_path | |
| if self.mode == 'test': | |
| pass | |
| else: | |
| label = annotation[1] | |
| imdb_['label'] = int(label) | |
| imdb_['flipped'] = False | |
| imdb_['bbox_target'] = np.zeros((4,)) | |
| imdb_['landmark_target'] = np.zeros((10,)) | |
| if len(annotation[2:])==4: | |
| bbox_target = annotation[2:6] | |
| imdb_['bbox_target'] = np.array(bbox_target).astype(float) | |
| if len(annotation[2:])==14: | |
| bbox_target = annotation[2:6] | |
| imdb_['bbox_target'] = np.array(bbox_target).astype(float) | |
| landmark = annotation[6:] | |
| imdb_['landmark_target'] = np.array(landmark).astype(float) | |
| imdb.append(imdb_) | |
| return imdb | |
| def append_flipped_images(self, imdb): | |
| """append flipped images to imdb | |
| Parameters: | |
| ---------- | |
| imdb: imdb | |
| image database | |
| Returns: | |
| ------- | |
| imdb: dict | |
| image database with flipped image annotations added | |
| """ | |
| print('append flipped images to imdb', len(imdb)) | |
| for i in range(len(imdb)): | |
| imdb_ = imdb[i] | |
| m_bbox = imdb_['bbox_target'].copy() | |
| m_bbox[0], m_bbox[2] = -m_bbox[2], -m_bbox[0] | |
| landmark_ = imdb_['landmark_target'].copy() | |
| landmark_ = landmark_.reshape((5, 2)) | |
| landmark_ = np.asarray([(1 - x, y) for (x, y) in landmark_]) | |
| landmark_[[0, 1]] = landmark_[[1, 0]] | |
| landmark_[[3, 4]] = landmark_[[4, 3]] | |
| item = {'image': imdb_['image'], | |
| 'label': imdb_['label'], | |
| 'bbox_target': m_bbox, | |
| 'landmark_target': landmark_.reshape((10)), | |
| 'flipped': True} | |
| imdb.append(item) | |
| self.image_set_index *= 2 | |
| return imdb | |
| 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']) | |
| if imdb[i]['flipped']: | |
| im = im[:, ::-1, :] | |
| 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 |