Update prediction.py

prediction.py

@@ -26,42 +26,20 @@ import torchvision.transforms as transforms
 from torch.utils.data import DataLoader
 from datasets import load_dataset
 
-# Define the path to the dataset
-dataset_path = 'andrewsunanda/fast_food_image_classification'
-
 # Load the dataset from Hugging Face
-dataset = load_dataset(dataset_path)
-
+from datasets import load_dataset
+dataset = load_dataset("andrewsunanda/fast_food_image_classification")
+    
 # Define the batch size and image size
 batch_size = 256
 img_size = (64, 64)
-
+    
 # Define the paths to the train, validation, and test folders
+dataset_path = "andrewsunanda/fast_food_image_classification"
 train_path = os.path.join(dataset_path, 'Train')
 valid_path = os.path.join(dataset_path, 'Valid')
 test_path = os.path.join(dataset_path, 'Test')
 
-# Define the transforms for the dataset
-transform = transforms.Compose([
-    transforms.Resize(img_size),
-    transforms.ToTensor(),
-])
-
-# Load the training dataset
-train_dataset = dataset['train']
-train_dataset = train_dataset.map(lambda x: {'image': transform(x['image']), 'label': x['label']})
-train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
-
-# Load the validation dataset
-valid_dataset = dataset['validation']
-valid_dataset = valid_dataset.map(lambda x: {'image': transform(x['image']), 'label': x['label']})
-valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False)
-
-# Load the testing dataset
-test_dataset = dataset['test']
-test_dataset = test_dataset.map(lambda x: {'image': transform(x['image']), 'label': x['label']})
-test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
-
 # Create data generators for training, validation, and testing
 train_datagen = ImageDataGenerator(
     rescale=1./255,