22 lines
823 B
Python
22 lines
823 B
Python
|
import torchvision.transforms as transforms
|
||
|
|
from torchvision.datasets import ImageFolder
|
||
|
|
from torch.utils.data import ConcatDataset
|
||
|
|
|
||
|
|
# images have to be the same size for the algorithm to work
|
||
|
|
transform = transforms.Compose([
|
||
|
|
transforms.Resize((224, 224)), # Resize images to (224, 224)
|
||
|
|
transforms.ToTensor(), # Convert images to tensors, 0-255 to 0-1
|
||
|
|
# transforms.RandomHorizontalFlip(), # 0.5 chance to flip the image
|
||
|
|
transforms.Normalize([0.5,0.5,0.5], [0.5,0.5,0.5])
|
||
|
|
])
|
||
|
|
|
||
|
|
letters_path = './letters'
|
||
|
|
package_path = './package'
|
||
|
|
|
||
|
|
# # Load images from folders
|
||
|
|
letter_folder = ImageFolder(letters_path, transform=transform)
|
||
|
|
package_folder = ImageFolder(package_path, transform=transform)
|
||
|
|
|
||
|
|
# Combine the both datasets into a single dataset
|
||
|
|
combined_dataset = ConcatDataset([letter_folder, package_folder])
|