import csv
import numpy as np
from PIL import Image
import pickle
from torchvision import transforms

#Please download the CheXpert 1.0 original image bundle and extract it into this working directory.
# Ensure that the file paths in img_list.csv match the extracted files.

img_num_list=[]
with open('img_list.csv', newline='') as csvfile:
    reader = csv.reader(csvfile)
    for i, row in enumerate(reader):
        if i>0:
            img_num_list.append([row[0],row[1]])


def load_resize_and_convert_image(image_path, target_size):

    img = Image.open(image_path).convert('L') #'RGB'


    img_resized = img.resize(target_size, Image.ANTIALIAS)


    img_array = np.array(img_resized)
    transform = transforms.Compose([
        transforms.ToTensor(),
    ])
    img_tensor = transform(img_resized)

    return img_array/255, img_tensor


img_dict={}
for i, a in enumerate(img_num_list):
    array, tensor=load_resize_and_convert_image(a[1], (128,128))
    img_dict[a[0]]=array


output_file_path = 'total_img_128.pkl'
with open(output_file_path, 'wb') as output_file:
    pickle.dump(img_dict, output_file)