def __init__(self, params, args):
super(gan, self).__init__()
self.G = MModel(params, use_cuda=True)
self.D = Discriminator(params, bias=True)
self.vgg_loss = VGGPerceptualLoss()
self.L1_loss = nn.L1Loss()
if args.use_cuda:
self.G = self.G.cuda()
self.D = self.D.cuda()
self.vgg_loss = self.vgg_loss.cuda()
self.L1_loss = self.L1_loss.cuda()
if args.g_weight_dir:
self.G.load_state_dict(torch.load(args.g_weight_dir), strict=True)
if args.d_weight_dir:
self.D.load_state_dict(torch.load(args.d_weight_dir), strict=False)
self.optimizer_G = torch.optim.Adam(self.G.parameters(), lr=args.g_lr)
self.optimizer_D = torch.optim.Adam(self.D.parameters(), lr=args.d_lr)
self.save_dir = args.save_dir
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
self.d_update_freq = args.d_update_freq
self.save_freq = args.save_freq
self.writer = SummaryWriter('runs/' + args.save_dir)
self.use_cuda = args.use_cuda
self.start_epoch = args.start_epoch
def __init__(self, params, args):
super(gan, self).__init__()
self.G = MModel(params, use_cuda=True)
# self.Face_G = Face_U_Net(img_ch=6, output_ch=3)
# self.Face_G = U_Net(params, output_ch=3)
self.Face_G = UNetWithResnet50Encoder(in_ch=3 + 3 + 1, n_classes=3)
# self.Face_D = FaceDisc()
# self.Face_D = Discriminator(params)
self.Face_D = PatchFace(in_channels=3 + 3 + 1)
self.vgg_loss = VGGPerceptualLoss()
self.L1_loss = nn.L1Loss()
if args.use_cuda:
self.G = self.G.cuda()
self.Face_G = self.Face_G.cuda()
self.Face_D = self.Face_D.cuda()
self.vgg_loss = self.vgg_loss.cuda()
self.L1_loss = self.L1_loss.cuda()
print('loading g model', args.g_weight_dir)
self.G.load_state_dict(torch.load(args.g_weight_dir), strict=True)
self.G.eval()
if args.face_g_weight_dir:
print('loading face_g model', args.face_g_weight_dir)
self.Face_G.load_state_dict(torch.load(args.face_g_weight_dir))
if args.face_d_weight_dir:
print('loading face_d model', args.face_d_weight_dir)
self.Face_D.load_state_dict(torch.load(args.face_d_weight_dir))
self.optimizer_Face_G = torch.optim.Adam(self.Face_G.parameters(),
lr=args.face_g_lr)
self.optimizer_Face_D = torch.optim.Adam(self.Face_D.parameters(),
lr=args.face_d_lr)
self.save_dir = args.save_dir
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
self.d_update_freq = args.d_update_freq
self.save_freq = args.save_freq
self.writer = SummaryWriter('runs/' + args.save_dir)
self.use_cuda = args.use_cuda
self.start_epoch = args.start_epoch
ret += [os.path.join(root, file)]
return ret
if __name__ == '__main__':
from net import MModel
from torch.utils.data import DataLoader
from torchvision import transforms
from torch.utils.tensorboard import SummaryWriter
import time
params = get_general_params()
params['IMG_HEIGHT'] = 256
params['IMG_WIDTH'] = 256
params['posemap_downsample'] = 2
net = MModel(params, use_cuda=True)
net.load_state_dict(torch.load(
'/versa/kangliwei/motion_transfer/0424-gan/g_epoch_2000.pth'),
strict=True)
net = net.cuda()
net.eval()
# ds = mtdataset(params, 'Standing Yoga Poses for Hips - Day 10 - The 30 Days of Yoga Challenge', 1, 'CHARLEY HULL 4K UHD SLOW MOTION FACE ON DRIVER GOLF SWING_1')
ds = mtdataset(
params,
'Standing Yoga Poses for Hips - Day 10 - The 30 Days of Yoga Challenge',
1, 'Tennis Tip_ Proper Weight Transfer On Topspin Groundstrokes')
dl = DataLoader(ds, 1, False)
writer = SummaryWriter(log_dir='runs/0428-test')
save_dir = '0428-ours-test-1'
if not os.path.exists(save_dir):
os.makedirs(save_dir)
# weight_dir = args.weight_dir
mini = not args.complete
full_y = args.full_y
shuffle = not args.no_shuffle
mini = False
full_y = False
shuffle = True
params = get_general_params()
params['IMG_HEIGHT'] = 256
params['IMG_WIDTH'] = 256
params['posemap_downsample'] = 2
ds = mtdataset(params, mini=mini, full_y=full_y)
dl = DataLoader(ds, bs, shuffle)
model = MModel(params, use_cuda=True).cuda()
model.train()
# model.load_state_dict(torch.load('/versa/kangliwei/motion_transfer/0424-gan/g_epoch_1670.pth'), strict=False)
# model.load_state_dict(torch.load('/versa/kangliwei/motion_transfer/0429-256-gan/g_epoch_720.pth'))
model.load_state_dict(torch.load('/versa/kangliwei/motion_transfer/0604-gan/epoch_290.pth'))
# if start_epoch != 0:
# model.load_state_dict(torch.load('/versa/kangliwei/motion_transfer/'+weight_dir+'/epoch_%d.pth'%(start_epoch-1)))
vgg_loss = VGGPerceptualLoss().cuda()
l1mask_loss = L1MaskLoss().cuda()
model_params = model.parameters()
optimizer = torch.optim.Adam(model_params, lr=lr)
writer = SummaryWriter('runs/'+model_dir)
print('len(dl)', len(dl))
print('len(ds)', len(ds))
for epoch in range(start_epoch, n_epoch):