Spaces:
Running
Running
| import torch | |
| import torchvision as tv | |
| import numpy as np | |
| import glob | |
| import os | |
| import random | |
| import cv2 | |
| from skimage import io | |
| from torch.utils.data import Dataset | |
| class LandmarkDataset(Dataset): | |
| def __init__(self, data_folder): | |
| super(LandmarkDataset, self).__init__() | |
| self.eids = sorted(os.listdir(data_folder)) | |
| self.images_folders = [os.path.join(data_folder, eid) for eid in self.eids] | |
| self.camera_ids = ['01', '02', '04', '05'] | |
| def get_item(self): | |
| landmarks = [] | |
| extrinsics = [] | |
| intrinsics = [] | |
| for images_folder in self.images_folders: | |
| landmarks_ = [] | |
| extrinsics_ = [] | |
| intrinsics_ = [] | |
| for v in range(len(self.camera_ids)): | |
| if os.path.exists(images_folder + '/landmarks_%s.npy' % self.camera_ids[v]): | |
| landmark = np.load(images_folder + '/landmarks_%s.npy' % self.camera_ids[v]) | |
| landmark = np.vstack([landmark[0:48], landmark[49:54], landmark[55:68]]) | |
| extrinsic = np.load(images_folder + '/camera_%s.npz' % self.camera_ids[v])['extrinsic'] | |
| intrinsic = np.load(images_folder + '/camera_%s.npz' % self.camera_ids[v])['intrinsic'] | |
| else: | |
| landmark = np.zeros([68, 3], dtype=np.float32) | |
| extrinsic = np.ones([3, 4], dtype=np.float32) | |
| intrinsic = np.ones([3, 3], dtype=np.float32) | |
| landmarks_.append(landmark) | |
| extrinsics_.append(extrinsic) | |
| intrinsics_.append(intrinsic) | |
| landmarks_ = np.stack(landmarks_) | |
| extrinsics_ = np.stack(extrinsics_) | |
| intrinsics_ = np.stack(intrinsics_) | |
| landmarks.append(landmarks_) | |
| extrinsics.append(extrinsics_) | |
| intrinsics.append(intrinsics_) | |
| landmarks = np.stack(landmarks) | |
| extrinsics = np.stack(extrinsics) | |
| intrinsics = np.stack(intrinsics) | |
| return landmarks, extrinsics, intrinsics, self.eids | |
| def __len__(self): | |
| return len(self.images_folders) | |