Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| from torch.utils.data import DataLoader | |
| from torchvision import transforms | |
| from datasets import load_dataset | |
| import torch.nn.functional as F | |
| from utils.config import load_config | |
| config = load_config() | |
| batch_size = config["batch_size"] | |
| num_workers = config["num_workers"] | |
| mean_nm = config["normalize_mean"] | |
| std_nm = config["normalize_std"] | |
| # let's select the first GPU device if available | |
| device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') | |
| dataset = load_dataset("DScomp380/plant_village") | |
| #split dataset into train(70%), and 30% remaining for val and test | |
| splits = dataset["train"].train_test_split(test_size=0.30, seed=42) | |
| train_split = splits["train"] #training set | |
| remaining = splits["test"] | |
| #split remaining 30% into val(15%) and test(15%) | |
| val_test = remaining.train_test_split(test_size=0.5, seed=42) | |
| val_split = val_test["train"] #validation set | |
| test_split = val_test["test"] #test set | |
| transform = transforms.Compose([ | |
| # resize images to 224x224, convert to tensor, and normalize | |
| transforms.Resize((224, 224)), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=mean_nm, std=std_nm) | |
| ]) | |
| def transform_batch(batch): | |
| # apply transformations to a batch of images | |
| batch["pixel_values"] = [transform(img) for img in batch["image"]] | |
| return batch | |
| # apply transformations to the datasets | |
| train_split = train_split.with_transform(transform_batch) | |
| val_split = val_split.with_transform(transform_batch) | |
| test_split = test_split.with_transform(transform_batch) | |
| def collate_fn(batch): | |
| # custom collate function to handle batching | |
| return { | |
| "pixel_values": torch.stack([item["pixel_values"] for item in batch]), | |
| "labels": torch.tensor([item["label"] for item in batch]) | |
| } | |
| # create DataLoaders for train, val, and test sets | |
| train_loader = DataLoader(train_split, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=collate_fn) | |
| val_loader = DataLoader(val_split, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate_fn) | |
| test_loader = DataLoader(test_split, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate_fn) | |
| if __name__ == "__main__": | |
| print(device) | |
| print(f"Loaded PlantVillage dataset with splits: {dataset}") |