'Image/Fake_frame',
fake_ending.permute(0, 2, 1, 3, 4).reshape(-1, C, H, W), epoch)
writer.add_video('Video/Input_video_fake', fake_video, epoch, fps=2)
writer.add_video('Video/Input_video_real', real_video, epoch, fps=2)
train_folder = './data/train/'
test_folder = './data/test/'
if __name__ == '__main__':
writer = SummaryWriter()
# Training dataset
train_dataset = data.VideoFolderDataset(train_folder,
cache=os.path.join(
train_folder, 'train.db'))
train_video_dataset = data.VideoDataset(train_dataset, 11)
train_loader = DataLoader(train_video_dataset,
batch_size=10,
drop_last=True,
num_workers=6,
shuffle=True)
test_dataset = data.VideoFolderDataset(test_folder,
cache=os.path.join(
test_folder, 'test.db'))
test_video_dataset = data.VideoDataset(test_dataset, 11)
test_loader = DataLoader(test_video_dataset,
batch_size=6,
drop_last=True,
num_workers=3,
shuffle=True)
##
if (chnum_in_ == 1):
norm_mean = [0.5]
norm_std = [0.5]
elif (chnum_in_ == 3):
norm_mean = (0.5, 0.5, 0.5)
norm_std = (0.5, 0.5, 0.5)
frame_trans = transforms.Compose(
[transforms.ToTensor(),
transforms.Normalize(norm_mean, norm_std)])
unorm_trans = utils.UnNormalize(mean=norm_mean, std=norm_std)
###### data
video_dataset = data.VideoDataset(tr_data_idx_dir,
tr_data_frame_dir,
transform=frame_trans)
tr_data_loader = DataLoader(video_dataset,
batch_size=batch_size_in,
shuffle=True,
num_workers=opt.NumWorker)
###### model
if (opt.ModelName == 'MemAE'):
model = AutoEncoderCov3DMem(chnum_in_,
mem_dim_in,
shrink_thres=sparse_shrink_thres)
else:
model = []
print('Wrong model name.')
model.apply(utils.weights_init)
dim_z_content = int(args['--dim_z_content']) # number of element in vector z_content
dim_z_motion = int(args['--dim_z_motion']) #number of element in vector z_motion
dim_z_category = int(args['--dim_z_category']) # number of element in vector z_category
dim_z_view = int(args['--dim_z_view'])
# Dataset of all
dataset = data.VideoFolderDataset(args['<dataset>'], cache=os.path.join(args['<dataset>'], 'local.db'))
# Object to get images from dataset VideoFolderDataset above
image_dataset = data.ImageDataset(dataset, image_transforms)
# Dataloader to load images
image_loader = DataLoader(image_dataset, batch_size=image_batch, drop_last=True, num_workers=2, shuffle=True)
# Object to get videos from dataset VideoFolderDataset above
video_dataset = data.VideoDataset(dataset, video_length=16, every_nth=2, transform=video_transforms)
# Dataloader to load videos
video_loader = DataLoader(video_dataset, batch_size=video_batch, drop_last=True, num_workers=2, shuffle=True)
# Create object of VideoGenerator
generator = models.VideoGenerator(n_channels, dim_z_content=dim_z_content, dim_z_view=dim_z_view,
dim_z_motion=dim_z_motion, dim_z_category=dim_z_category,
video_length=video_length)
# Create object of ImageDiscriminator
image_discriminator = build_discriminator(args['--image_discriminator'], n_channels=n_channels,
dim_z_view=dim_z_view, use_noise=args['--use_noise'],
noise_sigma=float(args['--noise_sigma']))
# Create object of VideoDiscriminator
video_discriminator = build_discriminator(args['--video_discriminator'], n_channels=n_channels,
video_batch = int(args['--video_batch'])
dim_z_content = int(args['--dim_z_content'])
dim_z_motion = int(args['--dim_z_motion'])
dim_z_category = int(args['--dim_z_category'])
# dataset = data.VideoFolderDataset(args['<dataset>'], cache=os.path.join(args['<dataset>'], 'local.db'))
dataset = data.VideoFolderDataset(args['<dataset>'], cache=None)
image_dataset = data.ImageDataset(dataset, image_transforms)
image_loader = DataLoader(image_dataset,
batch_size=image_batch,
drop_last=True,
num_workers=2,
shuffle=True)
video_dataset = data.VideoDataset(dataset, 16, 2, video_transforms)
video_loader = DataLoader(video_dataset,
batch_size=video_batch,
drop_last=True,
num_workers=2,
shuffle=True)
generator = models.VideoGenerator(n_channels, dim_z_content,
dim_z_category, dim_z_motion,
video_length)
image_discriminator = build_discriminator(args['--image_discriminator'],
n_channels=n_channels,
use_noise=args['--use_noise'],
noise_sigma=float(
args['--noise_sigma']))
if not os.path.exists(out_path):
os.makedirs(out_path)
for i, vid in enumerate(new_videos):
out = cv2.VideoWriter(f'{out_path}/vid_{i}.avi',
cv2.VideoWriter_fourcc(*"MPEG"), fps,
(320, 200))
for frame in vid:
frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
frame *= 255
frame = frame.astype(np.uint8)
out.write(frame)
cv2.imshow("frame", frame)
cv2.waitKey(100)
out.release()
sleep(2)
test_folder = './data/test/'
if __name__ == '__main__':
test_dataset = data.VideoFolderDataset(test_folder,
cache=os.path.join(
test_folder, 'test.db'))
test_video_dataset = data.VideoDataset(test_dataset, 11)
test_loader = DataLoader(test_video_dataset,
batch_size=5,
drop_last=True,
num_workers=3,
shuffle=True)
generate_video(test_loader)