Function main

tokenize_examples/tokenize_co3d_muse.py:72–136

(argv)

def main(argv):
    assert FLAGS.input_dir != ''
    assert FLAGS.output_file != ''

    # Load the pre-trained vq model from the hub
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    net = VQGANModel.from_pretrained('/home/vqlm/muse/ckpts/laion').to(device)
    net.eval()

    dirs = []
    for d in list_dir_with_full_path(FLAGS.input_dir):
        if not os.path.isdir(d):
            continue
        for d2 in list_dir_with_full_path(d):
            if not os.path.isdir(d2):
                continue
            dirs.append(d2)

    image_dirs = []
    for d in dirs:
        image_dirs.append((
            os.path.join(d, 'images'),
            os.path.join(d, 'masks'),
            os.path.join(d, 'depth_masks')
        ))

    with open(FLAGS.output_file, 'w') as fout:
        with torch.no_grad():
            for _ in trange(FLAGS.n_epochs, ncols=0):
                print(image_dirs[0])
                dataset = Co3DDataset(image_dirs, FLAGS.n_frames)
                dataloader = torch.utils.data.DataLoader(
                    dataset,
                    batch_size=FLAGS.batch_size * FLAGS.n_shots,
                    shuffle=False,
                    num_workers=FLAGS.n_workers,
                    drop_last=True
                )

                for batch in tqdm(dataloader, ncols=0):
                    batch_shape = batch.shape[:-3]
                    batch = batch.reshape(-1, *batch.shape[-3:])
                    batch = batch.permute(0,3,1,2)
                    batch = batch.to(device)

                    _, tokens = net.encode(batch)
                    tokens = tokens.reshape(*batch_shape, tokens.shape[-1])
                    # batch x task x frame x token
                    tokens = einops.rearrange(
                        tokens.cpu().numpy().astype(np.int32), '(b s) t f d -> b s t f d',
                        s=FLAGS.n_shots
                    )

                    image_mask_tokens = np.concatenate(
                        (tokens[:, :, 0, :, :], tokens[:, :, 1, :, :]), axis=-2
                    )
                    image_depth_tokens = np.concatenate(