mirror of
https://github.com/jantic/DeOldify.git
synced 2026-09-01 15:36:32 +08:00
Jason Antic
2a300aa84f
Getting rid of usused modules; Putting in mising type hints; Deleting usused logic.
243 lines
10 KiB
Python
243 lines
10 KiB
Python
from numpy import ndarray
|
|
from fastai.torch_imports import *
|
|
from fastai.core import *
|
|
from matplotlib.axes import Axes
|
|
from fastai.dataset import FilesDataset, ImageData, ModelData, open_image
|
|
from fastai.transforms import Transform, scale_min, tfms_from_stats, inception_stats
|
|
from fastai.transforms import CropType, NoCrop
|
|
from fasterai.training import GenResult, CriticResult, GANTrainer
|
|
from fasterai.images import ModelImageSet, EasyTensorImage
|
|
from IPython.display import display
|
|
from tensorboardX import SummaryWriter
|
|
import torchvision.utils as vutils
|
|
import statistics
|
|
|
|
|
|
class ModelImageVisualizer():
|
|
def __init__(self, default_sz:int=500):
|
|
self.default_sz=default_sz
|
|
|
|
def plot_transformed_image(self, path:Path, model:nn.Module, ds:FilesDataset, figsize:(int,int)=(20,20), sz:int=None,
|
|
tfms:[Transform]=[], compare:bool=True):
|
|
result = self.get_transformed_image_ndarray(path, model,ds, sz, tfms=tfms)
|
|
if compare:
|
|
orig = open_image(str(path))
|
|
fig,axes = plt.subplots(1, 2, figsize=figsize)
|
|
self.plot_image_from_ndarray(orig, axes=axes[0], figsize=figsize)
|
|
self.plot_image_from_ndarray(result, axes=axes[1], figsize=figsize)
|
|
else:
|
|
self.plot_image_from_ndarray(result, figsize=figsize)
|
|
|
|
def get_transformed_image_ndarray(self, path:Path, model:nn.Module, ds:FilesDataset, sz:int=None, tfms:[Transform]=[]):
|
|
training = model.training
|
|
model.eval()
|
|
orig = self.get_model_ready_image_ndarray(path, model, ds, sz, tfms)
|
|
orig = VV(orig[None])
|
|
result = model(orig).detach().cpu().numpy()
|
|
result = ds.denorm(result)
|
|
if training:
|
|
model.train()
|
|
return result[0]
|
|
|
|
def _transform(self, orig:ndarray, tfms:[Transform], model:nn.Module, sz:int):
|
|
for tfm in tfms:
|
|
orig,_=tfm(orig, False)
|
|
_,val_tfms = tfms_from_stats(inception_stats, sz, crop_type=CropType.NO, aug_tfms=[])
|
|
val_tfms.tfms = [tfm for tfm in val_tfms.tfms if not isinstance(tfm, NoCrop)]
|
|
orig = val_tfms(orig)
|
|
return orig
|
|
|
|
def get_model_ready_image_ndarray(self, path:Path, model:nn.Module, ds:FilesDataset, sz:int=None, tfms:[Transform]=[]):
|
|
im = open_image(str(path))
|
|
sz = self.default_sz if sz is None else sz
|
|
im = scale_min(im, sz)
|
|
im = self._transform(im, tfms, model, sz)
|
|
return im
|
|
|
|
def plot_image_from_ndarray(self, image:ndarray, axes:Axes=None, figsize=(20,20)):
|
|
if axes is None:
|
|
_,axes = plt.subplots(figsize=figsize)
|
|
clipped_image =np.clip(image,0,1)
|
|
axes.imshow(clipped_image)
|
|
axes.axis('off')
|
|
|
|
|
|
def plot_images_from_image_sets(self, image_sets:[ModelImageSet], validation:bool, figsize:(int,int)=(20,20),
|
|
max_columns:int=6, immediate_display:bool=True):
|
|
num_sets = len(image_sets)
|
|
num_images = num_sets * 2
|
|
rows, columns = self._get_num_rows_columns(num_images, max_columns)
|
|
|
|
fig, axes = plt.subplots(rows, columns, figsize=figsize)
|
|
title = 'Validation' if validation else 'Training'
|
|
fig.suptitle(title, fontsize=16)
|
|
|
|
for i, image_set in enumerate(image_sets):
|
|
self.plot_image_from_ndarray(image_set.orig.array, axes=axes.flat[i*2])
|
|
self.plot_image_from_ndarray(image_set.gen.array, axes=axes.flat[i*2+1])
|
|
|
|
if immediate_display:
|
|
display(fig)
|
|
|
|
|
|
def plot_image_outputs_from_model(self, ds:FilesDataset, model:nn.Module, idxs:[int], figsize:(int,int)=(20,20), max_columns:int=6,
|
|
immediate_display:bool=True):
|
|
image_sets = ModelImageSet.get_list_from_model(ds=ds, model=model, idxs=idxs)
|
|
self.plot_images_from_image_sets(image_sets=image_sets, figsize=figsize, max_columns=max_columns, immediate_display=immediate_display)
|
|
|
|
def _get_num_rows_columns(self, num_images:int, max_columns:int):
|
|
columns = min(num_images, max_columns)
|
|
rows = num_images//columns
|
|
rows = rows if rows * columns == num_images else rows + 1
|
|
return rows, columns
|
|
|
|
|
|
class ModelGraphVisualizer():
|
|
def __init__(self):
|
|
return
|
|
|
|
def write_model_graph_to_tensorboard(self, ds:FilesDataset, model:nn.Module, tbwriter:SummaryWriter):
|
|
try:
|
|
x,_=ds[0]
|
|
tbwriter.add_graph(model, V(x[None]))
|
|
except Exception as e:
|
|
print(("Failed to generate graph for model: {0}. Note that there's an outstanding issue with "
|
|
+ "scopes being addressed here: https://github.com/pytorch/pytorch/pull/12400").format(e))
|
|
|
|
class ModelHistogramVisualizer():
|
|
def __init__(self):
|
|
return
|
|
|
|
def write_tensorboard_histograms(self, model:nn.Module, iter_count:int, tbwriter:SummaryWriter):
|
|
for name, param in model.named_parameters():
|
|
tbwriter.add_histogram('/weights/' + name, param, iter_count)
|
|
|
|
|
|
|
|
class ModelStatsVisualizer():
|
|
def __init__(self):
|
|
return
|
|
|
|
def write_tensorboard_stats(self, model:nn.Module, iter_count:int, tbwriter:SummaryWriter):
|
|
gradients = [x.grad for x in model.parameters() if x.grad is not None]
|
|
gradient_nps = [to_np(x.data) for x in gradients]
|
|
|
|
if len(gradients) == 0:
|
|
return
|
|
|
|
avg_norm = sum(x.data.norm() for x in gradients)/len(gradients)
|
|
tbwriter.add_scalar('/gradients/avg_norm', avg_norm, iter_count)
|
|
|
|
median_norm = statistics.median(x.data.norm() for x in gradients)
|
|
tbwriter.add_scalar('/gradients/median_norm', median_norm, iter_count)
|
|
|
|
max_norm = max(x.data.norm() for x in gradients)
|
|
tbwriter.add_scalar('/gradients/max_norm', max_norm, iter_count)
|
|
|
|
min_norm = min(x.data.norm() for x in gradients)
|
|
tbwriter.add_scalar('/gradients/min_norm', min_norm, iter_count)
|
|
|
|
num_zeros = sum((np.asarray(x)==0.0).sum() for x in gradient_nps)
|
|
tbwriter.add_scalar('/gradients/num_zeros', num_zeros, iter_count)
|
|
|
|
|
|
avg_gradient= sum(x.data.mean() for x in gradients)/len(gradients)
|
|
tbwriter.add_scalar('/gradients/avg_gradient', avg_gradient, iter_count)
|
|
|
|
median_gradient = statistics.median(x.data.median() for x in gradients)
|
|
tbwriter.add_scalar('/gradients/median_gradient', median_gradient, iter_count)
|
|
|
|
max_gradient = max(x.data.max() for x in gradients)
|
|
tbwriter.add_scalar('/gradients/max_gradient', max_gradient, iter_count)
|
|
|
|
min_gradient = min(x.data.min() for x in gradients)
|
|
tbwriter.add_scalar('/gradients/min_gradient', min_gradient, iter_count)
|
|
|
|
class ImageGenVisualizer():
|
|
def __init__(self):
|
|
self.model_vis = ModelImageVisualizer()
|
|
|
|
def output_image_gen_visuals(self, md:ImageData, model:nn.Module, iter_count:int, tbwriter:SummaryWriter, jupyter:bool=False):
|
|
self._output_visuals(ds=md.val_ds, model=model, iter_count=iter_count, tbwriter=tbwriter, jupyter=jupyter, validation=True)
|
|
self._output_visuals(ds=md.trn_ds, model=model, iter_count=iter_count, tbwriter=tbwriter, jupyter=jupyter, validation=False)
|
|
|
|
def _output_visuals(self, ds:FilesDataset, model:nn.Module, iter_count:int, tbwriter:SummaryWriter,
|
|
validation:bool, jupyter:bool=False):
|
|
#TODO: Parameterize these
|
|
start_idx=0
|
|
count = 8
|
|
end_index = start_idx + count
|
|
idxs = list(range(start_idx,end_index))
|
|
image_sets = ModelImageSet.get_list_from_model(ds=ds, model=model, idxs=idxs)
|
|
self._write_tensorboard_images(image_sets=image_sets, iter_count=iter_count, tbwriter=tbwriter, validation=validation)
|
|
if jupyter:
|
|
self._show_images_in_jupyter(image_sets, validation=validation)
|
|
|
|
def _write_tensorboard_images(self, image_sets:[ModelImageSet], iter_count:int, tbwriter:SummaryWriter, validation:bool):
|
|
orig_images = []
|
|
gen_images = []
|
|
real_images = []
|
|
|
|
for image_set in image_sets:
|
|
orig_images.append(image_set.orig.tensor)
|
|
gen_images.append(image_set.gen.tensor)
|
|
real_images.append(image_set.real.tensor)
|
|
|
|
prefix = 'val' if validation else 'train'
|
|
|
|
tbwriter.add_image(prefix + ' orig images', vutils.make_grid(orig_images, normalize=True), iter_count)
|
|
tbwriter.add_image(prefix + ' gen images', vutils.make_grid(gen_images, normalize=True), iter_count)
|
|
tbwriter.add_image(prefix + ' real images', vutils.make_grid(real_images, normalize=True), iter_count)
|
|
|
|
|
|
def _show_images_in_jupyter(self, image_sets:[ModelImageSet], validation:bool):
|
|
#TODO: Parameterize these
|
|
figsize=(20,20)
|
|
max_columns=4
|
|
immediate_display=True
|
|
self.model_vis.plot_images_from_image_sets(image_sets, figsize=figsize, max_columns=max_columns,
|
|
immediate_display=immediate_display, validation=validation)
|
|
|
|
|
|
class GANTrainerStatsVisualizer():
|
|
def __init__(self):
|
|
return
|
|
|
|
def write_tensorboard_stats(self, gresult:GenResult, cresult:CriticResult, iter_count:int, tbwriter:SummaryWriter):
|
|
tbwriter.add_scalar('/loss/hingeloss', cresult.hingeloss, iter_count)
|
|
tbwriter.add_scalar('/loss/dfake', cresult.dfake, iter_count)
|
|
tbwriter.add_scalar('/loss/dreal', cresult.dreal, iter_count)
|
|
tbwriter.add_scalar('/loss/gcost', gresult.gcost, iter_count)
|
|
tbwriter.add_scalar('/loss/gcount', gresult.iters, iter_count)
|
|
tbwriter.add_scalar('/loss/gaddlloss', gresult.gaddlloss, iter_count)
|
|
|
|
def print_stats_in_jupyter(self, gresult:GenResult, cresult:CriticResult):
|
|
print(f'\nHingeLoss {cresult.hingeloss}; RScore {cresult.dreal}; FScore {cresult.dfake}; GAddlLoss {gresult.gaddlloss}; ' +
|
|
f'Iters: {gresult.iters}; GCost: {gresult.gcost};')
|
|
|
|
|
|
class LearnerStatsVisualizer():
|
|
def __init__(self):
|
|
return
|
|
|
|
def write_tensorboard_stats(self, metrics, iter_count:int, tbwriter:SummaryWriter):
|
|
if isinstance(metrics, list):
|
|
tbwriter.add_scalar('/loss/trn_loss', metrics[0], iter_count)
|
|
if len(metrics) == 1: return
|
|
tbwriter.add_scalar('/loss/val_loss', metrics[1], iter_count)
|
|
if len(metrics) == 2: return
|
|
|
|
for metric in metrics[2:]:
|
|
name = metric.__name__
|
|
tbwriter.add_scalar('/loss/'+name, metric, iter_count)
|
|
|
|
else:
|
|
tbwriter.add_scalar('/loss/trn_loss', metrics, iter_count)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|