Spaces:
Runtime error
Runtime error
File size: 2,799 Bytes
cb8043e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import torch
import cv2
import numpy as np
import torchvision.transforms as transforms
from torch.utils.data import Dataset
class ImageDataset(Dataset):
def __init__(self, csv, train, test):
self.csv = csv
self.train = train
self.test = test
self.all_image_names = self.csv[:]['Id']
self.all_labels = np.array(self.csv.drop(['Id', 'Genre'], axis=1))
self.train_ratio = int(0.85 * len(self.csv))
self.valid_ratio = len(self.csv) - self.train_ratio
# set the training data images and labels
if self.train == True:
print(f"Number of training images: {self.train_ratio}")
self.image_names = list(self.all_image_names[:self.train_ratio])
self.labels = list(self.all_labels[:self.train_ratio])
# define the training transforms
self.transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((400, 400)),
transforms.RandomHorizontalFlip(p=0.5),
transforms.RandomRotation(degrees=45),
transforms.ToTensor(),
])
# set the validation data images and labels
elif self.train == False and self.test == False:
print(f"Number of validation images: {self.valid_ratio}")
self.image_names = list(self.all_image_names[-self.valid_ratio:-10])
self.labels = list(self.all_labels[-self.valid_ratio:])
# define the validation transforms
self.transform = transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((400, 400)),
transforms.ToTensor(),
])
# set the test data images and labels, only last 10 images
# this, we will use in a separate inference script
elif self.test == True and self.train == False:
self.image_names = list(self.all_image_names[-10:])
self.labels = list(self.all_labels[-10:])
# define the test transforms
self.transform = transforms.Compose([
transforms.ToPILImage(),
transforms.ToTensor(),
])
def __len__(self):
return len(self.image_names)
def __getitem__(self, index):
image = cv2.imread(f"../input/movie-classifier/Multi_Label_dataset/Images/{self.image_names[index]}.jpg")
# convert the image from BGR to RGB color format
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# apply image transforms
image = self.transform(image)
targets = self.labels[index]
return {
'image': torch.tensor(image, dtype=torch.float32),
'label': torch.tensor(targets, dtype=torch.float32)
} |