)
assert 0
print(args)
print(os.getcwd())
device = torch.device(args.cuda)
state_dict = torch.load(args.model_dir, map_location=args.cuda)
model = Unet().to(device)
model.load_state_dict(state_dict)
custom_dataset = CustomDataset(root_dir=args.dataset)
dataloader = DataLoader(custom_dataset, args.batch_size, shuffle=False)
os.makedirs(os.path.join(args.save_dir, 'img'), exist_ok=True)
os.makedirs(os.path.join(args.save_dir, 'mask'), exist_ok=True)
if args.logdir is not None:
writer = SummaryWriter(args.logdir)
model.eval()
for i, batch in enumerate(tqdm(dataloader)):
img = batch.to(device)
with torch.no_grad():
pred, loss = model(img)
pred = pred[5].data
pred.requires_grad_(False)
if args.logdir is not None:
writer.add_image(args.model_dir + ', img', img, i)
writer.add_image(args.model_dir + ', mask', pred, i)
if args.save_dir is not None:
for j in range(img.shape[0]):
torchvision.utils.save_image(
img[j],
os.path.join(args.save_dir, 'img',
'{}_{}.jpg'.format(i, j)))
sage_model = Unet().to(device)
print("Model loaded! Loading Checkpoint...")
bdda_model_name = models[0]
sage_model_name = models[1]
bdda_state_dict = torch.load(os.path.join(model_dir, bdda_model_name))
sage_state_dict = torch.load(os.path.join(model_dir, sage_model_name))
bdda_model.load_state_dict(bdda_state_dict)
sage_model.load_state_dict(sage_state_dict)
print("Checkpoint loaded! Now predicting...")
bdda_model.eval()
sage_model.eval()
print('==============================')
img_shape = (1280, 720)
demo_dir = 'demo_images'
img_name = '{}/demo_img.jpg'.format(demo_dir)
img = Image.open(img_name)
img = img.convert('RGB')
sample = {'image': img}
