One text2Image dataset is tested

jloveric · May 17, 2024 · f97d053 · f97d053
1 parent 66f112c
commit f97d053
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 7 deletions.
diff --git a/high_order_implicit_representation/single_image_dataset.py b/high_order_implicit_representation/single_image_dataset.py
@@ -408,7 +408,7 @@ def __getitem__(self, idx):
         # ans = self.dataset()
         # print('ans', ans)
 
-        return next(self.gen_data())
+        return self.gen_data()
 
 
 class Text2ImageDataModule(LightningDataModule):

diff --git a/tests/test_single_image_dataset.py b/tests/test_single_image_dataset.py
@@ -38,9 +38,7 @@ def test_image_neighborhood_reader():
 
 
 def test_parquet_dataset():
-    dataset = Text2ImageDataset(
-        filenames=["test_data/test.parquet"]
-    )
+    dataset = Text2ImageDataset(filenames=["test_data/test.parquet"])
     dataloader = DataLoader(dataset, batch_size=2, shuffle=False)
 
     # caption, position, rgb = next(iter(dataloader))
@@ -51,6 +49,14 @@ def test_parquet_dataset():
 
 
 def test_text_to_image_sampler_dataloader():
-    dataloader = Text2ImageRenderDataset(
-        filenames=["test_data/test.parquet"]
-    )
+    dataset = Text2ImageRenderDataset(filenames=["test_data/test.parquet"])
+
+    dataloader = DataLoader(dataset, batch_size=2, shuffle=False)
+
+    for element in dataloader:
+        print("element", element[0][0].shape)
+        print("element", element[0][1].shape)
+        print("element", element[1][0].shape)
+        assert element[0].shape[0] ==2
+        assert len(element) == 3
+        break