import os,sys

import numpy as np

from PIL import Image as PILImage

import torch

import torch.nn.functional as F

from torch.utils import data as data

from torchvision import transforms as transforms

# dataset for sign detection and char detection

class COVID_CT_DATA(data.Dataset):

     def __init__(self, **kwargs):           

         super(COVID_CT_DATA).__init__()

         self.stage = kwargs['stage']

         # this returns the path to data dir

         self.data = kwargs['data']         

         self.fs = sorted(os.listdir(self.data))

         self.size = kwargs['img_size']

         # this returns the path to 

         self.img_fname = None

     def transform_img(self, img):

         # Faster R-CNN does the normalization

         t_ = transforms.Compose([

                             #transforms.ToPILImage(),

                             transforms.Resize(self.size),

                             transforms.ToTensor(),

                             ])

         img = t_(img)

         return img

     def load_img_label(self, idx):

         lab=torch.zeros(3, dtype=torch.float)

         lab[int(self.fs[idx].split('_')[0])] = 1

         im = PILImage.open(os.path.join(self.data, self.fs[idx]))

         if im.mode !='RGB':

            im = im.convert(mode='RGB')

         im = self.transform_img(im)

         return im, lab

     #'magic' method: size of the dataset

     def __len__(self):

         return len(os.listdir(self.data))

     # return one datapoint

     def __getitem__(self, idx):

         X,y = self.load_img_label(idx)

         return X,y