|
|
|
|
|
|
|
""" |
|
@Author : Peike Li |
|
@Contact : [email protected] |
|
@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 |
|
|
