import csv
import json
import os

import datasets
import pandas as pd
import numpy as np

ds = datasets.load_dataset('./coco_dataset_script.py', data_dir='/home/33611/caption/')
ds = ds['train']


def transform(example):

    example['pixel_values'] = np.load(example['pixels_file'])
    return example


# ds = ds.map(transform)

n = 0
for x in ds:
    n += 1
    assert os.path.isfile(x['image_file'])
    if n == 10:
        print(x)
        break