from data.robotic_contour_dataloader import get_train_dataloader, get_val_dataloader
from utils.options import Options
from data.utils.prepare_data import get_split
from train_engine_contour import TrainEngine
if __name__ == '__main__':
opt = Options().opt
train_files, test_files = get_split(opt.fold)
train_dataloader = get_train_dataloader(train_files, opt)
val_dataloader = get_val_dataloader(test_files, opt)
engine = TrainEngine(opt)
engine.set_data(train_dataloader, val_dataloader)
engine.train_model()
num_workers=opt.workers,
pin_memory=True)
return train_dataloader
def get_val_dataloader(file_list, opt):
data_transform = Compose([
PadIfNeeded(
min_height=opt.val_crop_height, min_width=opt.val_crop_width, p=1),
CenterCrop(height=opt.val_crop_height, width=opt.val_crop_width, p=1),
Normalize(p=1)
],
p=1)
val_dataset = RoboticsDataset(file_names=file_list,
transform=data_transform,
problem_type=opt.problem_type)
val_dataloader = torch.utils.data.DataLoader(val_dataset,
batch_size=opt.batch_size,
shuffle=False,
num_workers=opt.workers,
pin_memory=True)
return val_dataloader
if __name__ == '__main__':
# Test code for dataloader
options = Options()
train_files, test_files = prepare_data.get_split(0)
get_train_dataloader(train_files, options.opt)