| import numpy as np |
| from torch.utils.data import DataLoader |
| from .PIFuDataset import PIFuDataset |
| import pytorch_lightning as pl |
|
|
|
|
| class PIFuDataModule(pl.LightningDataModule): |
|
|
| def __init__(self, cfg): |
| super(PIFuDataModule, self).__init__() |
| self.cfg = cfg |
| self.overfit = self.cfg.overfit |
|
|
| if self.overfit: |
| self.batch_size = 1 |
| else: |
| self.batch_size = self.cfg.batch_size |
|
|
| self.data_size = {} |
|
|
| def prepare_data(self): |
|
|
| pass |
|
|
| @staticmethod |
| def worker_init_fn(worker_id): |
| np.random.seed(np.random.get_state()[1][0] + worker_id) |
|
|
| def setup(self, stage): |
|
|
| if stage == 'fit': |
| self.train_dataset = PIFuDataset(cfg=self.cfg, split="train") |
| self.val_dataset = PIFuDataset(cfg=self.cfg, split="val") |
| self.data_size = { |
| 'train': len(self.train_dataset), |
| 'val': len(self.val_dataset) |
| } |
|
|
| if stage == 'test': |
| self.test_dataset = PIFuDataset(cfg=self.cfg, split="test") |
|
|
| def train_dataloader(self): |
|
|
| train_data_loader = DataLoader(self.train_dataset, |
| batch_size=self.batch_size, |
| shuffle=True, |
| num_workers=self.cfg.num_threads, |
| pin_memory=True, |
| worker_init_fn=self.worker_init_fn) |
|
|
| return train_data_loader |
|
|
| def val_dataloader(self): |
|
|
| if self.overfit: |
| current_dataset = self.train_dataset |
| else: |
| current_dataset = self.val_dataset |
|
|
| val_data_loader = DataLoader(current_dataset, |
| batch_size=1, |
| shuffle=False, |
| num_workers=self.cfg.num_threads, |
| pin_memory=True, |
| worker_init_fn=self.worker_init_fn) |
|
|
| return val_data_loader |
|
|
| def test_dataloader(self): |
|
|
| test_data_loader = DataLoader(self.test_dataset, |
| batch_size=1, |
| shuffle=False, |
| num_workers=self.cfg.num_threads, |
| pin_memory=True) |
|
|
| return test_data_loader |
|
|
