Update prediction.py

prediction.py

@@ -20,15 +20,47 @@ def preprocess_input_image(img_path):
     x /= 255.
     return x, img1
 
-main_path= 'D:\\tugas_andrew_DS\\phase_2\\m2\\food'
-    
-# Define batch size and image size
+import os
+import torch
+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)
+
+# Define the batch size and image size
 batch_size = 256
 img_size = (64, 64)
-# Define paths to the data folders
-train_path = os.path.join(main_path, 'Train')
-valid_path = os.path.join(main_path, 'Valid')
-test_path = os.path.join(main_path, 'Test')
+
+# Define the paths to the train, validation, and test folders
+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(