|
import matplotlib |
|
|
|
matplotlib.use("Agg") |
|
|
|
import os, sys |
|
import yaml |
|
from argparse import ArgumentParser |
|
from time import gmtime, strftime |
|
from shutil import copy |
|
|
|
from frames_dataset import FramesDataset |
|
|
|
from modules.generator import OcclusionAwareGenerator, OcclusionAwareSPADEGenerator |
|
from modules.discriminator import MultiScaleDiscriminator |
|
from modules.keypoint_detector import KPDetector, HEEstimator |
|
|
|
import torch |
|
|
|
from train import train |
|
|
|
if __name__ == "__main__": |
|
if sys.version_info[0] < 3: |
|
raise Exception("You must use Python 3 or higher. Recommended version is Python 3.7") |
|
|
|
parser = ArgumentParser() |
|
parser.add_argument("--config", default="config/vox-256.yaml", help="path to config") |
|
parser.add_argument( |
|
"--mode", |
|
default="train", |
|
choices=[ |
|
"train", |
|
], |
|
) |
|
parser.add_argument("--gen", default="original", choices=["original", "spade"]) |
|
parser.add_argument("--log_dir", default="log", help="path to log into") |
|
parser.add_argument("--checkpoint", default=None, help="path to checkpoint to restore") |
|
parser.add_argument( |
|
"--device_ids", |
|
default="0, 1, 2, 3, 4, 5, 6, 7", |
|
type=lambda x: list(map(int, x.split(","))), |
|
help="Names of the devices comma separated.", |
|
) |
|
parser.add_argument("--verbose", dest="verbose", action="store_true", help="Print model architecture") |
|
parser.set_defaults(verbose=False) |
|
|
|
opt = parser.parse_args() |
|
with open(opt.config) as f: |
|
config = yaml.load(f, Loader=yaml.FullLoader) |
|
|
|
if opt.checkpoint is not None: |
|
log_dir = os.path.join(*os.path.split(opt.checkpoint)[:-1]) |
|
else: |
|
log_dir = os.path.join(opt.log_dir, os.path.basename(opt.config).split(".")[0]) |
|
log_dir += " " + strftime("%d_%m_%y_%H.%M.%S", gmtime()) |
|
|
|
if opt.gen == "original": |
|
generator = OcclusionAwareGenerator(**config["model_params"]["generator_params"], **config["model_params"]["common_params"]) |
|
elif opt.gen == "spade": |
|
generator = OcclusionAwareSPADEGenerator(**config["model_params"]["generator_params"], **config["model_params"]["common_params"]) |
|
|
|
if torch.cuda.is_available(): |
|
print("cuda is available") |
|
generator.to(opt.device_ids[0]) |
|
if opt.verbose: |
|
print(generator) |
|
|
|
discriminator = MultiScaleDiscriminator(**config["model_params"]["discriminator_params"], **config["model_params"]["common_params"]) |
|
if torch.cuda.is_available(): |
|
discriminator.to(opt.device_ids[0]) |
|
if opt.verbose: |
|
print(discriminator) |
|
|
|
kp_detector = KPDetector(**config["model_params"]["kp_detector_params"], **config["model_params"]["common_params"]) |
|
|
|
if torch.cuda.is_available(): |
|
kp_detector.to(opt.device_ids[0]) |
|
|
|
if opt.verbose: |
|
print(kp_detector) |
|
|
|
he_estimator = HEEstimator(**config["model_params"]["he_estimator_params"], **config["model_params"]["common_params"]) |
|
|
|
if torch.cuda.is_available(): |
|
he_estimator.to(opt.device_ids[0]) |
|
|
|
dataset = FramesDataset(is_train=(opt.mode == "train"), **config["dataset_params"]) |
|
|
|
if not os.path.exists(log_dir): |
|
os.makedirs(log_dir) |
|
if not os.path.exists(os.path.join(log_dir, os.path.basename(opt.config))): |
|
copy(opt.config, log_dir) |
|
|
|
if opt.mode == "train": |
|
print("Training...") |
|
train(config, generator, discriminator, kp_detector, he_estimator, opt.checkpoint, log_dir, dataset, opt.device_ids) |
|
|