if os.path.isfile(config.dataset_path):
pathfile = open(config.dataset_path, 'rt').read().splitlines()
elif os.path.isdir(config.dataset_path):
pathfile = glob.glob(os.path.join(config.dataset_path, '*.png'))
else:
print('Invalid testing data file/folder path.')
exit(1)
total_number = len(pathfile)
test_num = total_number if config.test_num == -1 else min(
total_number, config.test_num)
print('The total number of testing images is {}, and we take {} for test.'.
format(total_number, test_num))
print('configuring model..')
ourModel = InpaintingModel_GMCNN(in_channels=4, 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')))
print('Loading done.')
if config.random_mask:
np.random.seed(config.seed)
for i in range(test_num):
if config.mask_path:
mask = get_file_mask(config.mask_path)
elif config.mask_type == 'rect':
mask, _ = generate_rect_mask(config.img_shapes, config.mask_shapes,
#========================================
print('loading data..')
dataset = InpaintingDataset(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(in_channels=4, opt=config)
ourModel.print_networks()
if config.load_model_dir != '':
print('Loading pretrained model from {}'.format(config.load_model_dir))
load_model_path = getLatest(os.path.join(config.load_model_dir, '*.pth'))
print("Load model path:", load_model_path)
ourModel.load_networks(load_model_path)
print('Loading done.')
# ourModel = torch.nn.DataParallel(ourModel).cuda()
print('model setting up..')
print('training initializing..')
writer = SummaryWriter(log_dir=tensorboard_exp_dir)
cnt = 0
for epoch in range(config.epochs):
for i, data in enumerate(dataloader):
config = TrainOptions().parse()
print('loading data..')
dataset = InpaintingDataset(config.dataset_path,
'',
transform=transforms.Compose([ToTensor()]))
dataloader = DataLoader(dataset,
batch_size=config.batch_size,
shuffle=True,
num_workers=0,
drop_last=True)
print('data loaded..')
print('Preparing model..')
ourModel = InpaintingModel_GMCNN(in_channels=4, opt=config)
if config.load_model_dir != '':
ourModel.load_networks(
getLatest(os.path.join(config.load_model_dir, '*.pth')))
print('Initializing training..')
for epoch in range(config.epochs):
for i, data in enumerate(dataloader):
gt = data['gt'].cuda()
gt = gt / 127.5 - 1
data_in = {'gt': gt}
ourModel.setInput(data_in)
import numpy as np
import cv2
import os
from options.test_options import TestOptions
from model.net import InpaintingModel_GMCNN
from util.utils import generate_my_mask, getLatest
config = TestOptions().parse()
ourModel = InpaintingModel_GMCNN(in_channels=4, opt=config)
if config.load_model_dir != '':
ourModel.load_networks(
getLatest(os.path.join(config.load_model_dir, '*.pth')))
def generate(pathfile, x1, y1, x2, y2, i):
mask, _ = generate_my_mask(config.img_shapes, x1, y1, x2 - x1, y2 - y1)
image = cv2.imread(pathfile)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = np.transpose(image, [2, 0, 1])
image = np.expand_dims(image, axis=0)
h, w = image.shape[2:]
grid = 4
image = image[:, :, :h // grid * grid, :w // grid * grid]
mask = mask[:, :, :h // grid * grid, :w // grid * grid]
result = ourModel.evaluate(image, mask)
result = np.transpose(result[0][::-1, :, :], [1, 2, 0])
cv2.imwrite(os.path.join(config.saving_path, '{:03d}.png'.format(i)),
result)