From 9d9863898002881668d1954bb1a946d95639eb08 Mon Sep 17 00:00:00 2001
From: Hunor Laczko <luni741@gmail.com>
Date: Fri, 9 Dec 2022 19:22:38 +0000
Subject: [PATCH] basic dataset

---
 .gitignore      |  1 +
 data/dataset.py | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+)
 create mode 100644 data/dataset.py

diff --git a/.gitignore b/.gitignore
index 9a86641..d4be5a2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -160,4 +160,5 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
+.vscode/
 temp/
\ No newline at end of file
diff --git a/data/dataset.py b/data/dataset.py
new file mode 100644
index 0000000..6153e18
--- /dev/null
+++ b/data/dataset.py
@@ -0,0 +1,37 @@
+import os
+import sys
+
+import math
+import torch
+import torchvision as tv
+import numpy as np
+
+from torch.utils.data import Dataset
+
+class Dataset(Dataset):
+    def __init__(self, folder):
+        self.folder = folder
+        self.transforms = tv.transforms.ToTensor()
+        self.frames = os.listdir(os.path.join(folder, 'mesh_posed'))
+
+    def __len__(self):
+        return len(self.frames)
+
+    def __getitem__(self, idx):
+        img_path = os.path.join(self.folder, 'mesh_posed', self.frames[idx])
+        image = np.load(img_path)['mesh_posed']
+        # image = self.transforms(image)
+        return image
+
+
+if __name__ == '__main__':
+    from torch.utils.data import DataLoader
+
+    dataset = Dataset('/home/workspace/local-data/train')
+    train_dataloader = DataLoader(dataset, batch_size=4, drop_last=True, num_workers=2)
+
+    for batch in train_dataloader:
+        batch
+        pass
+
+    pass 
\ No newline at end of file