import tensorflow as tfGCS_PATH_FULL_RESOUTION = "gs://beans-fullres/tfrecords"
GCS_PATH_LOW_RESOLUTION = "gs://beans-lowres/tfrecords"BATCH_SIZE = 4
AUTO = tf.data.AUTOTUNEdef parse_tfr(proto):
    feature_description = {
        "image": tf.io.VarLenFeature(tf.float32),
        "image_shape": tf.io.VarLenFeature(tf.int64),
        "label": tf.io.VarLenFeature(tf.int64),
    }
    rec = tf.io.parse_single_example(proto, feature_description)
    image_shape = tf.sparse.to_dense(rec["image_shape"])
    image = tf.reshape(tf.sparse.to_dense(rec["image"]), image_shape)
    label = tf.sparse.to_dense(rec["label"])
    return {"pixel_values": image, "label": label}


def prepare_dataset(GCS_PATH=GCS_PATH_FULL_RESOUTION, 
                    split="train", batch_size=BATCH_SIZE):

    if split not in ["train", "val"]:
        raise ValueError(
            "Invalid split provided. Supports splits are: `train` and `val`."
        )

    dataset = tf.data.TFRecordDataset(
        [filename for filename in tf.io.gfile.glob(f"{GCS_PATH}/{split}-*")],
        num_parallel_reads=AUTO,
    ).map(parse_tfr, num_parallel_calls=AUTO)

    if split == "train":
        dataset = dataset.shuffle(batch_size * 2)

    dataset = dataset.batch(batch_size)
    dataset = dataset.prefetch(AUTO)
    return datasettrain_dataset = prepare_dataset()
val_dataset = prepare_dataset(split="val")for batch in train_dataset.take(1):
    print(batch["pixel_values"].shape, batch["label"].shape)for batch in val_dataset.take(1):
    print(batch["pixel_values"].shape, batch["label"].shape)train_dataset = prepare_dataset(GCS_PATH_LOW_RESOLUTION)
val_dataset = prepare_dataset(GCS_PATH_LOW_RESOLUTION, split="val")for batch in train_dataset.take(1):
    print(batch["pixel_values"].shape, batch["label"].shape)for batch in val_dataset.take(1):
    print(batch["pixel_values"].shape, batch["label"].shape)