Datatrooper/posters_classification

1import torch2import cv23import numpy as np4import torchvision.transforms as transforms5from torch.utils.data import Dataset6 7 8class ImageDataset(Dataset):9    def __init__(self, csv, train, test):10        self.csv = csv11        self.train = train12        self.test = test13        self.all_image_names = self.csv[:]['Id']14        self.all_labels = np.array(self.csv.drop(['Id', 'Genre'], axis=1))15        self.train_ratio = int(0.85 * len(self.csv))16        self.valid_ratio = len(self.csv) - self.train_ratio17        # set the training data images and labels18        if self.train == True:19            print(f"Number of training images: {self.train_ratio}")20            self.image_names = list(self.all_image_names[:self.train_ratio])21            self.labels = list(self.all_labels[:self.train_ratio])22            # define the training transforms23            self.transform = transforms.Compose([24                transforms.ToPILImage(),25                transforms.Resize((400, 400)),26                transforms.RandomHorizontalFlip(p=0.5),27                transforms.RandomRotation(degrees=45),28                transforms.ToTensor(),29            ])30        # set the validation data images and labels31        elif self.train == False and self.test == False:32            print(f"Number of validation images: {self.valid_ratio}")33            self.image_names = list(self.all_image_names[-self.valid_ratio:-10])34            self.labels = list(self.all_labels[-self.valid_ratio:])35            # define the validation transforms36            self.transform = transforms.Compose([37                transforms.ToPILImage(),38                transforms.Resize((400, 400)),39                transforms.ToTensor(),40            ])41        # set the test data images and labels, only last 10 images42        # this, we will use in a separate inference script43        elif self.test == True and self.train == False:44            self.image_names = list(self.all_image_names[-10:])45            self.labels = list(self.all_labels[-10:])46             # define the test transforms47            self.transform = transforms.Compose([48                transforms.ToPILImage(),49                transforms.ToTensor(),50            ])51    def __len__(self):52        return len(self.image_names)53    54    def __getitem__(self, index):55        image = cv2.imread(f"../input/movie-classifier/Multi_Label_dataset/Images/{self.image_names[index]}.jpg")56        # convert the image from BGR to RGB color format57        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)58        # apply image transforms59        image = self.transform(image)60        targets = self.labels[index]61        62        return {63            'image': torch.tensor(image, dtype=torch.float32),64            'label': torch.tensor(targets, dtype=torch.float32)65        }