def _get_vectorised_colour_samples(archetype_config, dataloader):
    (num_batches, batch_sz) = (len(dataloader), archetype_config.dataloader_batch_sz)
    num_imgs_max = (num_batches * batch_sz)
    img_sz = archetype_config.input_sz
    imgs_all = np.zeros((num_imgs_max, img_sz, img_sz, archetype_config.in_channels), dtype=np.uint8)
    masks_all = np.zeros((num_imgs_max, img_sz, img_sz), dtype=np.bool)
    labels_all = None
    actual_num_imgs = 0
    for (b_i, batch) in enumerate(dataloader):
        if (len(batch) == 3):
            store_labels = True
            if (labels_all is None):
                labels_all = np.zeros((num_imgs_max, img_sz, img_sz), dtype=np.int32)
            (imgs, labels, masks) = batch
            labels = labels.cpu().numpy().astype(np.int32)
        else:
            store_labels = False
            (imgs, _, _, masks) = batch
        imgs = (imgs * 255.0).permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
        masks = masks.cpu().numpy().astype(np.bool)
        (curr_batch_sz, h, w, c) = imgs.shape
        assert ((h == archetype_config.input_sz) and (w == archetype_config.input_sz) and (c == archetype_config.in_channels))
        if (b_i < (num_batches - 1)):
            assert (batch_sz == curr_batch_sz)
        start = (b_i * batch_sz)
        imgs_all[start:(start + curr_batch_sz), :, :, :] = imgs
        masks_all[start:(start + curr_batch_sz), :, :] = masks
        if store_labels:
            labels_all[start:(start + curr_batch_sz), :, :] = labels
        actual_num_imgs += curr_batch_sz
    imgs_all = imgs_all[:actual_num_imgs, :, :, :]
    masks_all = masks_all[:actual_num_imgs, :, :]
    num_unmasked = masks_all.sum()
    if store_labels:
        labels_all = labels_all[:actual_num_imgs, :, :]
        samples_labels = labels_all[masks_all].reshape((- 1))
        assert (samples_labels.shape[0] == num_unmasked)
    samples = imgs_all[masks_all, :].reshape((- 1), archetype_config.in_channels)
    assert (samples.shape[0] == num_unmasked)
    if (not store_labels):
        return samples
    else:
        return (samples, samples_labels)
def _get_vectorised_colour_samples(archetype_config, dataloader):
    (num_batches, batch_sz) = (len(dataloader), archetype_config.dataloader_batch_sz)
    num_imgs_max = (num_batches * batch_sz)
    img_sz = archetype_config.input_sz
    imgs_all = np.zeros((num_imgs_max, img_sz, img_sz, archetype_config.in_channels), dtype=np.uint8)
    masks_all = np.zeros((num_imgs_max, img_sz, img_sz), dtype=np.bool)
    labels_all = None
    actual_num_imgs = 0
    for (b_i, batch) in enumerate(dataloader):
        if (len(batch) == 3):
            store_labels = True
            if (labels_all is None):
                labels_all = np.zeros((num_imgs_max, img_sz, img_sz), dtype=np.int32)
            (imgs, labels, masks) = batch
            labels = labels.cpu().numpy().astype(np.int32)
        else:
            store_labels = False
            (imgs, _, _, masks) = batch
        imgs = (imgs * 255.0).permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)
        masks = masks.cpu().numpy().astype(np.bool)
        (curr_batch_sz, h, w, c) = imgs.shape
        assert ((h == archetype_config.input_sz) and (w == archetype_config.input_sz) and (c == archetype_config.in_channels))
        if (b_i < (num_batches - 1)):
            assert (batch_sz == curr_batch_sz)
        start = (b_i * batch_sz)
        imgs_all[start:(start + curr_batch_sz), :, :, :] = imgs
        masks_all[start:(start + curr_batch_sz), :, :] = masks
        if store_labels:
            labels_all[start:(start + curr_batch_sz), :, :] = labels
        actual_num_imgs += curr_batch_sz
    imgs_all = imgs_all[:actual_num_imgs, :, :, :]
    masks_all = masks_all[:actual_num_imgs, :, :]
    num_unmasked = masks_all.sum()
    if store_labels:
        labels_all = labels_all[:actual_num_imgs, :, :]
        samples_labels = labels_all[masks_all].reshape((- 1))
        assert (samples_labels.shape[0] == num_unmasked)
    samples = imgs_all[masks_all, :].reshape((- 1), archetype_config.in_channels)
    assert (samples.shape[0] == num_unmasked)
    if (not store_labels):
        return samples
    else:
        return (samples, samples_labels)