Working on new sampler

high_order_implicit_representation/rendering.py

@@ -11,7 +11,7 @@
 from high_order_implicit_representation.single_image_dataset import (
     image_neighborhood_dataset,
     image_to_dataset,
-    Text2ImageDataset
+    Text2ImageRenderDataset
 )
 import math
 import matplotlib.pyplot as plt
@@ -244,7 +244,7 @@ def on_train_epoch_end(
 
 class Text2ImageGenerator(Callback):
     def __init__(self, filename, rotations, batch_size):
-        self._dataset = Text2ImageDataset(filename, rotations=rotations)
+        self._dataset = Text2ImageRenderDataset(filename, rotations=rotations)
         self._dataloader = DataLoader(self._dataset, batch_size=batch_size, shuffle=False)
         self._batch_size = batch_size
 
@@ -258,12 +258,12 @@ def on_train_epoch_end(
 
             y_hat_list = []
 
-            for batch in self._dataloader:
-                res = pl_module(
-                    self._inputs[
-                        batch * self._batch_size : (batch + 1) * self._batch_size
-                    ]
-                )
+            for caption_embedding, flattened_image, flattened_position in self._dataloader:
+                for index, rgb in enumerate(flattened_image):
+                    res = pl_module(
+                        caption_embedding, flattened_position[index], rgb
+                    )
+                    
                 y_hat_list.append(res.detach().cpu())
             y_hat = torch.cat(y_hat_list)
 

high_order_implicit_representation/single_image_dataset.py

@@ -380,6 +380,37 @@ def __getitem__(self, idx):
         return next(self.gen_data())
 
 
+class Text2ImageRenderDataset(Dataset):
+    """
+    This one is just used for drawing complete images, I want the entire
+    image returned as a flattened list, not just pixels
+    """
+    def __init__(self, filenames: List[str]):
+        super().__init__()
+        self.dataset = PickAPic(files=filenames)
+        self.sentence_model = SentenceTransformer("all-MiniLM-L6-v2")
+
+    def __len__(self):
+        return int(1e12)
+
+    def gen_data(self):
+
+        caption, image = next(self.dataset())
+        caption_embedding = self.sentence_model.encode(caption)
+
+        print("got here")
+        flattened_image, flattened_position, image = simple_image_to_dataset(image)
+        return caption_embedding, flattened_image, flattened_position
+
+
+    def __getitem__(self, idx):
+        # I'm totally ignoring the index
+        # ans = self.dataset()
+        # print('ans', ans)
+
+        return next(self.gen_data())
+
+
 class Text2ImageDataModule(LightningDataModule):
     def __init__(self, filenames: List[str], num_workers, batch_size, pin_memory):
         super().__init__()