from torchvision import transforms
import torchvision.utils as vutils
from tensorboardX import SummaryWriter
from data.data import InpaintingDataset, ToTensor
from model.net import InpaintingModel_DFBM
from options.train_options import TrainOptions
from util.utils import getLatest
from multiprocessing import freeze_support
if __name__ == '__main__':
config = TrainOptions().parse()
print('loading data..')
dataset = InpaintingDataset(config.data_file,
config.dataset_path,
transform=transforms.Compose([ToTensor()]))
dataloader = DataLoader(dataset,
batch_size=config.batch_size,
shuffle=True,
num_workers=4,
drop_last=True)
print('data loaded..')
print('configuring model..')
ourModel = InpaintingModel_DFBM(opt=config)
ourModel.print_networks()
if config.load_model_dir != '':
print('Loading pretrained model from {}'.format(config.load_model_dir))
ourModel.load_networks(
getLatest(os.path.join(config.load_model_dir, '*.pth')))
py_file_name = config.py_file.split("/")[
-1] # Get python file name (soruce code name)
checkpoint_dir = os.path.join(config.out_dir, py_file_name + "/checkpoints")
os.makedirs(checkpoint_dir, exist_ok=True)
# make tensorboard subdirectory for the experiment
tensorboard_exp_dir = os.path.join(config.tensorboard_dir, py_file_name)
os.makedirs(tensorboard_exp_dir, exist_ok=True)
#========================================
print('loading data..')
dataset = InpaintingDataset_WithMask_v2(config.dataset_path,
config.data_root,
transform=transforms.Compose(
[ToTensor()]))
dataloader = DataLoader(dataset,
batch_size=config.batch_size,
shuffle=True,
num_workers=4,
drop_last=True)
print('data loaded..')
print('configuring model..')
ourModel = InpaintingModel_GMCNN_Given_Mask(in_channels=4, opt=config)
ourModel.print_networks()
if (config.load_model_dir != '') or (config.load_model_path != ''):
print('Loading pretrained model from {}'.format(config.load_model_dir))
#ourModel.load_networks(getLatest(os.path.join(config.load_model_dir, '*.pth')))
# Modified by Vajira to load exact path
ourModel.load_networks(config.load_model_path)