Untitled

mail@pastecode.io avatar
unknown
python
2 years ago
1.2 kB
1
Indexable
Never
# image normalization
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# preprocessing of images
class IIITDataset(Dataset):
    def __init__(self, image_paths, transform):
        super().__init__()
        self.paths = image_paths
        self.len = len(self.paths)
        self.transform = transform

    def __len__(self): return self.len

    def __getitem__(self, index): 
        path = self.paths[index]
        image = Image.open(path).convert('RGB')
        image = self.transform(image)
        label = int(df['class_id'].loc[df['image']==path])
        #label = 0 if 'Abyssinian' in path else 1
        
        return (image, label)


# create train dataset
import random

# create train-test split
random.shuffle(img_files)

train = img_files[:3680]
test = img_files[3680:]

print("train size", len(train))
print("test size", len(test))



train_ds = IIITDataset(train, transform)
train_dl = DataLoader(train_ds, batch_size=100)
print(len(train_ds), len(train_dl))

# create test dataset
test_ds = IIITDataset(test, transform)
test_dl = DataLoader(test_ds, batch_size=100)
print(len(test_ds), len(test_dl))