File size: 556 Bytes
e857f97 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | import cv2
import torch
from data.build_datasets import DataInfo
from data.process_thermal import get_thermal_transform
from torchvision import datasets
def get_thermal_dataset(args):
data_path = args.thermal_data_path
transform = get_thermal_transform(args)
dataset = datasets.ImageFolder(data_path, transform=transform)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
num_workers=args.workers,
sampler=None,
)
return DataInfo(dataloader=dataloader, sampler=None)
|