[AI] Covid, Pneumonia 데이터 실습

Covid, Pneumonia 폐 사진을 가지고 분류하는 실습을 해보자.
둘 다 걸려봐서 다 내 사진인줄

Dataset 구성

class CustomDataset(Dataset):
    def __init__(self, data_path, mode, transform=None):
        """ init : 초기값 설정 """
        self.all_data = sorted(
            glob.glob(os.path.join(data_path, mode, "*", "*")))
        self.transform = transform

    def __getitem__(self, index):
        data_path = self.all_data[index]
        # data_path info >>  ./Covid19-dataset/train/Covid/01.jpeg
        print("data_path info >> ", data_path)
        data_path_split = data_path.split("/")
        print(data_path_split[3])
        labels_temp = data_path_split[3]

        label = -1
        if "Covid" == labels_temp:
            label = 0
        elif "Normal" == labels_temp:
            label = 1
        elif "Viral Pneumonia" == labels_temp:
            label = 3

        image = Image.open(data_path).convert('RGB')

        if self.transform is not None:
            image = self.transform(image)
        print(image, label)
        return image, label

    def __len__(self):
        return len(self.all_data)


image_transform = transforms.Compose([
    transforms.ToTensor()
])

train_data = CustomDataset(
    "./Covid19-dataset/", "train", transform=image_transform)
test_data = CustomDataset("./Covid19-dataset/", "test",
                          transform=image_transform)