Function prepare_data

script/dm/prepare_data.py:10–35

prepare data for ready to train posenet, return dataloaders

(args, images, poses_train, i_split)

to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8)

def prepare_data(args, images, poses_train, i_split):
    ''' prepare data for ready to train posenet, return dataloaders '''
    #TODO: Convert GPU friendly data generator later: https://stanford.edu/~shervine/blog/pytorch-how-to-generate-data-parallel
    #TODO: Probably a better implementation style here: https://github.com/PyTorchLightning/pytorch-lightning
    
    i_train, i_val, i_test = i_split

    img_train = torch.Tensor(images[i_train]).permute(0, 3, 1, 2) # now shape is [N, CH, H, W]
    pose_train = torch.Tensor(poses_train[i_train])
    
    trainset = TensorDataset(img_train, pose_train)
    train_dl = DataLoader(trainset, batch_size=args.batch_size, shuffle=True)
    
    img_val = torch.Tensor(images[i_val]).permute(0, 3, 1, 2) # now shape is [N, CH, H, W]
    pose_val = torch.Tensor(poses_train[i_val])
    
    valset = TensorDataset(img_val, pose_val)
    val_dl = DataLoader(valset)

    img_test = torch.Tensor(images[i_test]).permute(0, 3, 1, 2) # now shape is [N, CH, H, W]
    pose_test = torch.Tensor(poses_train[i_test])

    testset = TensorDataset(img_test, pose_test)
    test_dl = DataLoader(testset)

    return train_dl, val_dl, test_dl

def load_dataset(args):
    ''' load posenet training data '''