def write_imgs(model, dset, device, s_batch=1, model_name='model'):
"""
Writing collages to 'show_cases/predictions'.
A collage feature the original image the L channel and the predicted image in that order from left to right.
:param model: model to be used
:param dset: data set, containing images in 'show_cases/imgs'
:param device: GPU device
:param s_batch: size of batch
:param model_name: name of model
"""
data_iter = data.RAMDataSetIter(dset, s_batch=s_batch, shuffle=False)
with torch.no_grad():
model.to(device)
model.eval()
it_count = 0
prog_bar = tqdm.tqdm(total=len(dset), bar_format=f'{Fore.GREEN}{model_name} show case images {Fore.RESET} ' +
"{l_bar}%s{bar}%s{r_bar}" % (Fore.GREEN, Fore.RESET))
while True:
try:
batch, targets = data_iter.next()
except StopIteration:
break
batch = batch.to(device)
targets = targets.to(device)
y = torch.nn.functional.hardtanh(model(batch), min_val=0, max_val=1).cpu().data.numpy()
paths = dset.paths[it_count:it_count+len(y)]
for i, pred_ab in enumerate(y):
# lightness channel
L = batch[i].cpu().data.numpy()
bw = np.repeat(L*255., 3, axis=0).transpose(1, 2, 0).astype(np.uint8)
# true a, b channels
img_ab = targets[i].cpu().data.numpy()
# ture image
img_Lab = np.vstack([L, img_ab]) * 255.
img_Lab = img_Lab.transpose(1, 2, 0).astype(np.uint8)
img_BGR = cv2.cvtColor(img_Lab, cv2.COLOR_LAB2BGR)
# estimated image
pred_Lab = np.vstack([L, pred_ab]) * 255.
pred_Lab = pred_Lab.transpose(1, 2, 0).astype(np.uint8)
pred_BGR = cv2.cvtColor(pred_Lab, cv2.COLOR_LAB2BGR)
# summary: ture image, lightness, estimated image
summary = np.hstack([img_BGR, bw, pred_BGR])
# channels: true channels, estimated channels
channels = np.hstack([np.vstack(img_ab), np.vstack(pred_ab)]) * 255.
channels = channels.astype(np.uint8)
# saving summary and channels
cv2.imwrite(join(f'show_cases/predictions/{model_name}_I{it_count+i+1}.png'), summary)
cv2.imwrite(join(f'show_cases/predictions/{model_name}_I{it_count+i+1}_channels.png'), channels)
it_count += 1
prog_bar.update(len(y))
model.cpu()
batch.cpu()
targets.cpu()
prog_bar.close()
def __init__(self, model, dset, device, s_batch=32, model_name='model'):
"""
:param model: ConvNet model
:param dset: data set to use in training
:param device: GPU device
:param s_batch: batch size
:param model_name: name of model
"""
self.model = model
self.dset = dset
self.device = device
self.s_batch = s_batch
self.model_name = model_name
self.print_c = Fore.GREEN
if not os.path.isdir(join(f'logs/{self.model_name}')):
os.makedirs(join(f'logs/{self.model_name}'))
def plot_progress(self, error_series, plt_imgs, plt_imgs_t):
"""
This method generates an informative plot, displaying the loss over time along logged images.
:param error_series: loss values over time (/ steps)
:param plt_imgs: logged images
:param plt_imgs_t: time (/ step) at which respective images have been logged
:return:
"""
fig, ax = plt.subplots(figsize=(160, 40))
ax.plot(np.arange(1, 1 + len(error_series)), error_series, alpha=0.9)
for i, p_I in enumerate(plt_imgs):
imagebox = OffsetImage(p_I, zoom=1.)
x = plt_imgs_t[i]
ab = AnnotationBbox(imagebox, (x, error_series[x - 1]),
xybox=(0., 256.),
xycoords='data',
boxcoords="offset points",
arrowprops={'arrowstyle': '->'})
ax.add_artist(ab)
plt.savefig(join(f'logs/{self.model_name}/progress.png'),
bbox_inches='tight',
dpi=50)
plt.close(fig)
hist = np.mean(hists, axis=0)
return bounds, hist
if __name__ == '__main__':
# establishing GPU device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# evaluating the models
accs = {}
print(
f'\n{Fore.BLUE} evaluating all models stored in {Fore.WHITE}models/{Fore.RESET}'
)
for model_name in sorted(
[p for p in os.listdir('models') if os.path.isdir(f'models/{p}')]):
# loading the model and corresponding dataset
with open(join(f'models/{model_name}/kwargs.json'), 'r') as f_kwargs:
kwargs = json.load(f_kwargs)
state_dict = torch.load(f'models/{model_name}/state_dict.pth',
map_location=torch.device('cpu'))
dset = data.RAMDataSet(**kwargs['dset'])
model = nn.ConvNet(**kwargs['model'])
model.load_state_dict(state_dict)
# initializing the evaluator
evaluator = Evaluator(model, dset, device, model_name=model_name)
# storing and saving some discrete upper bounds, for a quick look/ preview
accs[model_name] = evaluator.evaluate()
print(
f'{Fore.BLUE}accuracy {model_name}: {Fore.WHITE}{accs[model_name]}{Fore.RESET}'
)
# constructing a plot visualizing the distribution of abs. differences, that lie within certain ranges
bounds, hist = evaluator.evaluate_hist()
def log_step(self, y, batch, targets, str_it_count):
"""
Saving a sample of the batch.
:param y: predicted ab channels
:param batch: true L channels
:param targets: true ab channels
:param str_it_count: prefix of the saved images (embodies the current optimization step)
:return:
"""
# index of sample for displayal
k = np.random.randint(0, y.shape[0])
# BW images
if y.shape[1] == 1:
img = y[k].cpu().numpy()[0]
img = np.repeat(img[:, :, None], 3, axis=2)
orig = targets[k].cpu().data.numpy()[0]
orig = np.repeat(orig[:, :, None], 3, axis=2)
# logging images
cv2.imwrite(join(f'logs/{self.model_name}/{str_it_count}_est.png'),
img)
cv2.imwrite(
join(f'logs/{self.model_name}/{str_it_count}_orig.png'), orig)
# assemble the plot iamge
img = cv2.resize(img, (128, 128))
orig = cv2.resize(orig, (128, 128))
p_I = np.vstack([orig, img])
# Lab images
else:
ab = y[k].cpu().data.numpy()
orig_ab = targets[k].cpu().data.numpy()
L = batch[k].cpu().data.numpy()
# estimated image and fmaps
est_Lab = np.vstack([L, ab]) * 255.
est_RGB = cv2.cvtColor(
est_Lab.transpose(1, 2, 0).astype(np.uint8),
cv2.COLOR_Lab2RGB) / 255.
est_Lab = np.hstack(est_Lab[1:]) / 255.
est_ab_BW = np.repeat(est_Lab[:, :, None], 3, axis=2)
# original image and fmaps
orig_Lab = np.vstack([L, orig_ab]) * 255.
orig_RGB = cv2.cvtColor(
orig_Lab.transpose(1, 2, 0).astype(np.uint8),
cv2.COLOR_Lab2RGB) / 255.
orig_Lab = np.hstack(orig_Lab[1:]) / 255.
orig_ab_BW = np.repeat(orig_Lab[:, :, None], 3, axis=2)
# logging images
cv2.imwrite(
join(f'logs/{self.model_name}/{str_it_count}_colour_est.png'),
cv2.cvtColor((est_RGB * 255).astype(np.uint8),
cv2.COLOR_RGB2BGR))
cv2.imwrite(
join(f'logs/{self.model_name}/{str_it_count}_ab_est.png'),
(est_ab_BW * 255).astype(np.uint8))
cv2.imwrite(
join(f'logs/{self.model_name}/{str_it_count}_colour_orig.png'),
cv2.cvtColor((orig_RGB * 255).astype(np.uint8),
cv2.COLOR_RGB2BGR))
cv2.imwrite(
join(f'logs/{self.model_name}/{str_it_count}_ab_orig.png'),
(orig_ab_BW * 255).astype(np.uint8))
# assemble compressed visual summary (the plot image)
est_RGB = cv2.resize(est_RGB, (128, 128))
orig_RGB = cv2.resize(orig_RGB, (128, 128))
est_ab_BW = cv2.resize(est_ab_BW, (2 * 128, 128))
orig_ab_BW = cv2.resize(orig_ab_BW, (2 * 128, 128))
p_I = np.vstack([
np.hstack([orig_RGB, orig_ab_BW]),
np.hstack([est_RGB, est_ab_BW])
])
return p_I
return ap
if __name__ == '__main__':
# checking for GPU device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# args parsing
args = arg_parser().parse_args()
print_info(args)
w, h = args.w, args.h
s_batch = args.s_batch
iterations = args.iterations
# loading the dataset
if args.dset == 'daisy':
dset_kwargs = dict(
path=join('/media/ssd1/data_sets/flowers/daisy/train'),
width=w,
height=h,
black_white=False,
suffix='.jpg')
dset = data.RAMDataSet(**dset_kwargs)
dset_kwargs['path'] = join('/media/ssd1/data_sets/flowers/daisy/test')
c_out = 2
else:
raise NotImplementedError(f'dataset {args.dset} is unknown')
# training models with 0-3 pooling operations
for n_pooling in [0, 1, 2, 3][::-1]:
# keeping track of the data set- and model parameters
model_kwargs = dict(c_out=c_out,
img_w=w,
img_h=h,
# summary: ture image, lightness, estimated image
summary = np.hstack([img_BGR, bw, pred_BGR])
# channels: true channels, estimated channels
channels = np.hstack([np.vstack(img_ab), np.vstack(pred_ab)]) * 255.
channels = channels.astype(np.uint8)
# saving summary and channels
cv2.imwrite(join(f'show_cases/predictions/{model_name}_I{it_count+i+1}.png'), summary)
cv2.imwrite(join(f'show_cases/predictions/{model_name}_I{it_count+i+1}_channels.png'), channels)
it_count += 1
prog_bar.update(len(y))
model.cpu()
batch.cpu()
targets.cpu()
prog_bar.close()
if __name__ == '__main__':
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'\n{Fore.BLUE}generating showcase images{Fore.RESET}')
for model_name in sorted(os.listdir('models')):
state_dict = torch.load(join(f'models/{model_name}/state_dict.pth'),
map_location=torch.device('cpu'))
with open(join(f'models/{model_name}/kwargs.json'), 'r') as f_kwargs:
kwargs = json.load(f_kwargs)
model = nn.ConvNet(**kwargs['model'])
model.load_state_dict(state_dict)
kwargs['dset']['path'] = join('show_cases/imgs/')
kwargs['dset']['suffix'] = '.jpg'
dset = data.RAMDataSet(**kwargs['dset'])
write_imgs(model, dset, device, model_name=model_name)