text: import os
  import torch
  from torch.utils.data import Dataset
  import glob
  import numpy as np
  from lib.utils.covid_utils  import  read_txt
  from PIL import Image
  import torchvision.transforms as transforms
  
  
  class CovidCTDataset(Dataset):
      def __init__(self,mode, root_dir, txt_COVID, txt_NonCOVID, transform=None):
          """
          Args:
              txt_path (string): Path to the txt file with annotations.
              root_dir (string): Directory with all the images.
              transform (callable, optional): Optional transform to be applied
                  on a sample.
          File structure:
          - root_dir
              - CT_COVID
                  - img1.png
                  - img2.png
                  - ......
              - CT_NonCOVID
                  - img1.png
                  - img2.png
                  - ......
          """
          self.root_dir = root_dir
          self.txt_path = [txt_COVID, txt_NonCOVID]
          self.classes = ['CT_COVID', 'CT_NonCOVID']
          self.num_cls = len(self.classes)
          self.img_list = []
          self.full_volume = None
          self.affine = None
          for c in range(self.num_cls):
              cls_list = [[os.path.join(self.root_dir, self.classes[c], item), c] for item in read_txt(self.txt_path[c])]
              self.img_list += cls_list
  
          normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
          train_transformer = transforms.Compose([
              transforms.Resize(256),
              transforms.RandomResizedCrop((224), scale=(0.5, 1.0)),
              transforms.RandomHorizontalFlip(),
              transforms.ToTensor(),
              normalize
          ])
  
          val_transformer = transforms.Compose([
              transforms.Resize(224),
              transforms.CenterCrop(224),
              transforms.ToTensor(),
              normalize
          ])
          if(mode == 'train'):
  
              self.transform = train_transformer
  
          else:
              self.transform = val_transformer
          print('samples = ', len(self.img_list))
  
  
  
  
  
      def __len__(self):
          return len(self.img_list)
  
      def __getitem__(self, idx):
          if torch.is_tensor(idx):
              idx = idx.tolist()
  
          img_path = self.img_list[idx][0]
  
          image = Image.open(img_path).convert('RGB')
  
          if self.transform:
              image = self.transform(image)
  
          return image, torch.tensor(int(self.img_list[idx][1]), dtype=torch.long)