| | |
| | |
| |
|
| | """ |
| | @Author : Peike Li |
| | @Contact : peike.li@yahoo.com |
| | @File : dataset.py |
| | @Time : 8/30/19 9:12 PM |
| | @Desc : Dataset Definition |
| | @License : This source code is licensed under the license found in the |
| | LICENSE file in the root directory of this source tree. |
| | """ |
| |
|
| | import os |
| | import pdb |
| |
|
| | import cv2 |
| | import numpy as np |
| | from PIL import Image |
| | from torch.utils import data |
| | from utils.transforms import get_affine_transform |
| |
|
| |
|
| | class SimpleFolderDataset(data.Dataset): |
| | def __init__(self, root, input_size=[512, 512], transform=None): |
| | self.root = root |
| | self.input_size = input_size |
| | self.transform = transform |
| | self.aspect_ratio = input_size[1] * 1.0 / input_size[0] |
| | self.input_size = np.asarray(input_size) |
| | self.is_pil_image = False |
| | if isinstance(root, Image.Image): |
| | self.file_list = [root] |
| | self.is_pil_image = True |
| | elif os.path.isfile(root): |
| | self.file_list = [os.path.basename(root)] |
| | self.root = os.path.dirname(root) |
| | else: |
| | self.file_list = os.listdir(self.root) |
| |
|
| | def __len__(self): |
| | return len(self.file_list) |
| |
|
| | def _box2cs(self, box): |
| | x, y, w, h = box[:4] |
| | return self._xywh2cs(x, y, w, h) |
| |
|
| | def _xywh2cs(self, x, y, w, h): |
| | center = np.zeros((2), dtype=np.float32) |
| | center[0] = x + w * 0.5 |
| | center[1] = y + h * 0.5 |
| | if w > self.aspect_ratio * h: |
| | h = w * 1.0 / self.aspect_ratio |
| | elif w < self.aspect_ratio * h: |
| | w = h * self.aspect_ratio |
| | scale = np.array([w, h], dtype=np.float32) |
| | return center, scale |
| |
|
| | def __getitem__(self, index): |
| | if self.is_pil_image: |
| | img = np.asarray(self.file_list[index])[:, :, [2, 1, 0]] |
| | else: |
| | img_name = self.file_list[index] |
| | img_path = os.path.join(self.root, img_name) |
| | img = cv2.imread(img_path, cv2.IMREAD_COLOR) |
| | h, w, _ = img.shape |
| |
|
| | |
| | person_center, s = self._box2cs([0, 0, w - 1, h - 1]) |
| | r = 0 |
| | trans = get_affine_transform(person_center, s, r, self.input_size) |
| | input = cv2.warpAffine( |
| | img, |
| | trans, |
| | (int(self.input_size[1]), int(self.input_size[0])), |
| | flags=cv2.INTER_LINEAR, |
| | borderMode=cv2.BORDER_CONSTANT, |
| | borderValue=(0, 0, 0)) |
| |
|
| | input = self.transform(input) |
| | meta = { |
| | 'center': person_center, |
| | 'height': h, |
| | 'width': w, |
| | 'scale': s, |
| | 'rotation': r |
| | } |
| |
|
| | return input, meta |
| |
|
