def __next__(self):
"""Get next training batch/test example.
Returns:
batch: dict, has "pixels" and "rays".
"""
if self.split == "train":
return utils.shard(self.queue.get())
else:
return utils.to_device(self.queue.get())
def peek(self):
"""Peek at the next training batch or test example without dequeuing it.
Returns:
batch: dict, has "pixels" and "rays".
"""
x = self.queue.queue[0].copy() # Make a copy of the front of the queue.
if self.split == "train":
return utils.shard(x)
else:
return utils.to_device(x)
