DataLoaders.from_dsets separate dependent and independent variables

Hi, I am trying to create a model that takes a list of images and returns also a list of images. For example, let’s say the input is 11 images. The dependent variables (which should be fed to the model) are the first 10 images, but the 11th image is an independent variable. The output of the model is going to be n images, let’s say 5 images. Then the loss is a function taking the 11th input image and the 5 output images.

I’m new to DL and to fastai, so I’m trying to make a very light use of fastai. This is my (incomplete) code so far (and this is a colab):

import random
import PIL

def list_folders(path):
  p = Path(path)
  return [p/f for f in os.listdir(path) if os.path.isdir(p/f)]

def open_image(fname, size=224):
    img = PIL.Image.open(fname).convert('RGB')
    img = img.resize((size, size))
    t = torch.Tensor(np.array(img))
    return t.permute(2,0,1).float()/255.0

class SpacesDataset(torch.utils.data.Dataset):
  def __init__(self, scene_paths, is_valid=False):
    self.is_valid = is_valid
    self.scenes = [ReadScene(p) for p in scene_paths]
    if is_valid:
      self.calculated_length = len(self.scenes)*len(self.scenes[0])
    else:
      self.calculated_length = len(self.scenes)*len(self.scenes[0])*len(self.scenes[0][0])
      
  def __getitem__(self, i):
    if self.is_valid:
      camera_i = i % len(self.scenes[0])
      scene_i = (i - camera_i) / len(self.scenes[0])
      photos = [self.scenes[scene_i][camera_i][k] for k in range(3)]
    else:
      photos = self._draw()

    for photo in photos:
      photo.image = open_image(photo.image_path)

    return [photo.image for photo in photos]
  
  def __len__(self): return self.calculated_length
  
  def _draw(self):
    scene_i = random.choice(range(len(self.scenes)))
    camera_i = random.choice(range(len(self.scenes[scene_i])))
    photo_is = np.random.permutation(len(self.scenes[scene_i][camera_i]))[:3]
    return [self.scenes[scene_i][camera_i][i] for i in photo_is]

scenes = list_folders("spaces_dataset/data/800/")
idxs = np.random.permutation(range(len(scenes)))
cut = int(0.8 * len(scenes))

train_data = [ scenes[x] for x in idxs[:cut] ]
valid_data = [ scenes[x] for x in idxs[cut:] ]

train_ds = SpacesDataset(train_data)
valid_ds = SpacesDataset(valid_data, is_valid=True)

from fastai.data.core import DataLoaders
dls = DataLoaders.from_dsets(train_ds, valid_ds)
b = dls.one_batch()
# Use the GPU
dls = dls.cuda()

# learn = Learner(dls, model, loss_func=CrossEntropyLossFlat())

This code of course is incomplete, I don’t know how to continue. What I think I need to do is to be able to create a Dataloaders that correctly separates the dependent and independent variables, then I’ll create a loss function that uses the output of the model and the independent variables. Could you guys help me with the Dataloaders part?