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])