parser.add_argument('--loaders', type=int, default=4, help='number of parallel data loading processes') parser.add_argument('--batch_size', type=int, default=32) args = parser.parse_args() pipeline.initialize(args) ## load model : image --> reflectance x normals x depth x lighting model = models.Decomposer().cuda() ## get a data loader for train and val sets train_set = pipeline.IntrinsicDataset(args.data_path, args.train_sets, args.intrinsics, array=args.array, size_per_dataset=args.num_train) train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, num_workers=args.loaders, shuffle=True) val_set = pipeline.IntrinsicDataset(args.data_path, args.val_sets, args.intrinsics, array=args.array, size_per_dataset=args.num_val) val_loader = torch.utils.data.DataLoader(val_set, batch_size=args.batch_size, num_workers=args.loaders,
train_set = pipeline.ComposerDataset(args.data_path, args.unlabeled, args.labeled, unlabeled_array=args.unlabeled_array, labeled_array=args.labeled_array, size_per_dataset=args.set_size) train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size / 2, num_workers=args.loader_threads, shuffle=True) ## data loader for val set, which is completely labeled val_set = pipeline.IntrinsicDataset(args.data_path, args.val_sets, args.val_intrinsics, inds=range(0, args.num_val * args.val_offset, args.val_offset), array=args.unlabeled_array, size_per_dataset=args.set_size) val_loader = torch.utils.data.DataLoader(val_set, batch_size=args.batch_size, num_workers=args.loader_threads, shuffle=False) ## print out error plots after every epoch for every prediction logger = pipeline.Logger( ['recon', 'refl', 'depth', 'shape', 'lights', 'shading'], args.save_path) param_updater = pipeline.ParamUpdater(args.transfer) for epoch in range(args.num_epochs): print '<Main> Epoch {}'.format(epoch)
shape_loss = 0 lights_loss = 0 # preparing necessary models #decomp_path = "saved/decomposer/decomp_with_dropout/state_dropout.t7" decomp_path = "saved/decomposer/decomp_cl/state_cl.t7" decomposer = models.Decomposer().cuda() checkpoint = torch.load(decomp_path) #decomposer.load_state_dict(checkpoint['model_state_dict']) decomposer.load_state_dict(checkpoint) decomposer.train(mode=False) test_set = pipeline.IntrinsicDataset(args.data_path, args.test_sets, args.intrinsics, array=args.array, size_per_dataset=args.num_test, rel_path='../intrinsics-network') test_loader = torch.utils.data.DataLoader(test_set, batch_size=args.batch_size, num_workers=args.loaders, shuffle=False) images = [] for ind, tensors in enumerate(test_loader): tensors = [Variable(t.float().cuda(async=True)) for t in tensors] inp, mask, refl_targ, depth_targ, shape_targ, lights_targ = tensors refl_pred, depth_pred, shape_pred, lights_pred = decomposer.forward( inp, mask)
import sys, os, argparse, torch import models, pipeline val_set = pipeline.IntrinsicDataset('dataset/output/','car_tarin',['input', 'mask', 'albedo', 'depth', 'normals', 'lights'], size_per_dataset=1) print ("is",type(val_set)) val_loader = torch.utils.data.DataLoader(val_set, batch_size=32, num_workers=4, shuffle=False) shader = models.Shader().cuda() # pipeline.visualize_shader(shader, val_loader, "saved/a.png",save_raw=True) for epoch in range(4): save_path = os.path.join('saved/shader/', str(epoch) + '.png') val_losses = pipeline.visualize_decomposer(shader, val_loader, "saved/a.png", 1,save_raw=True)