diff --git a/ColorizeTrainingWide.ipynb b/ColorizeTrainingWide.ipynb
new file mode 100644
index 00000000..d65ded9f
--- /dev/null
+++ b/ColorizeTrainingWide.ipynb
@@ -0,0 +1,914 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Pretrained GAN"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import os\n",
+ "os.environ['CUDA_VISIBLE_DEVICES']='0' "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import fastai\n",
+ "from fastai import *\n",
+ "from fastai.vision import *\n",
+ "from fastai.callbacks.tensorboard import *\n",
+ "from fastai.vision.gan import *\n",
+ "from fasterai.generators import *\n",
+ "from fasterai.critics import *\n",
+ "from fasterai.dataset import *\n",
+ "from fasterai.loss import *\n",
+ "from PIL import Image, ImageDraw, ImageFont\n",
+ "from PIL import ImageFile"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Setup"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "path = Path('data/imagenet/ILSVRC/Data/CLS-LOC')\n",
+ "path_hr = path\n",
+ "path_lr = path/'bandw'\n",
+ "\n",
+ "proj_id = 'ColorizeNew73'\n",
+ "gen_name = proj_id + '_gen'\n",
+ "crit_name = proj_id + '_crit'\n",
+ "\n",
+ "name_gen = proj_id + '_image_gen'\n",
+ "path_gen = path/name_gen\n",
+ "\n",
+ "TENSORBOARD_PATH = Path('data/tensorboard/' + proj_id)\n",
+ "\n",
+ "nf_factor = 2"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def save_all(suffix=''):\n",
+ " learn_gen.save(gen_name + str(sz) + suffix)\n",
+ " learn_crit.save(crit_name + str(sz) + suffix)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def load_all(suffix=''):\n",
+ " learn_gen.load(gen_name + str(sz) + suffix, with_opt=False)\n",
+ " learn_crit.load(crit_name + str(sz) + suffix, with_opt=False)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def get_data(bs:int, sz:int, keep_pct:float):\n",
+ " return get_colorize_data(sz=sz, bs=bs, crappy_path=path_lr, good_path=path_hr, \n",
+ " random_seed=None, keep_pct=keep_pct)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def get_crit_data(classes, bs, sz):\n",
+ " src = ImageList.from_folder(path, include=classes, recurse=True).random_split_by_pct(0.1, seed=42)\n",
+ " ll = src.label_from_folder(classes=classes)\n",
+ " data = (ll.transform(get_transforms(max_zoom=2.), size=sz)\n",
+ " .databunch(bs=bs).normalize(imagenet_stats))\n",
+ " return data"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def crappify(fn,i):\n",
+ " dest = path_lr/fn.relative_to(path_hr)\n",
+ " dest.parent.mkdir(parents=True, exist_ok=True)\n",
+ " img = PIL.Image.open(fn).convert('LA').convert('RGB')\n",
+ " img.save(dest) "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def save_preds(dl):\n",
+ " i=0\n",
+ " names = dl.dataset.items\n",
+ " \n",
+ " for b in dl:\n",
+ " preds = learn_gen.pred_batch(batch=b, reconstruct=True)\n",
+ " for o in preds:\n",
+ " o.save(path_gen/names[i].name)\n",
+ " i += 1"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def save_gen_images(learn_gen):\n",
+ " if path_gen.exists(): shutil.rmtree(path_gen)\n",
+ " path_gen.mkdir(exist_ok=True)\n",
+ " data_gen = get_data(bs=bs, sz=sz, keep_pct=0.085)\n",
+ " save_preds(data_gen.fix_dl)\n",
+ " PIL.Image.open(path_gen.ls()[0])"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## Crappified data"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Prepare the input data by crappifying images."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Uncomment the first time you run this notebook."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "#il = ImageItemList.from_folder(path_hr)\n",
+ "#parallel(crappify, il.items)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Pre-training"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Pre-train generator"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Now let's pretrain the generator."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "bs=88\n",
+ "sz=64\n",
+ "keep_pct=1.0"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "data_gen = get_data(bs=bs, sz=sz, keep_pct=keep_pct)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen = gen_learner_deep(arch=models.resnet101, data=data_gen, gen_loss=FeatureLoss2(), nf_factor=nf_factor)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.callback_fns.append(partial(ImageGenTensorboardWriter, base_dir=TENSORBOARD_PATH, name='GenPre'))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.fit_one_cycle(2, pct_start=0.8, max_lr=slice(1e-3))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.save(gen_name)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.load(gen_name, with_opt=False)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.unfreeze()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.fit_one_cycle(2, pct_start=0.01, max_lr=slice(3e-7, 3e-4))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.save(gen_name)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "bs=20\n",
+ "sz=128\n",
+ "keep_pct=1.0"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.data = get_data(sz=sz, bs=bs, keep_pct=keep_pct)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.unfreeze()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.load(gen_name, with_opt=False)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.fit_one_cycle(2, pct_start=0.01, max_lr=slice(1e-7,1e-4))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.save(gen_name)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.load(gen_name, with_opt=False)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "bs=8\n",
+ "sz=192\n",
+ "keep_pct=0.50"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.data = get_data(sz=sz, bs=bs, keep_pct=keep_pct)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.unfreeze()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.fit_one_cycle(1, pct_start=0.01, max_lr=slice(5e-8,5e-5))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen.save(gen_name)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Save generated images"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "save_gen_images(gen_name)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Train critic"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Pretrain the critic on crappy vs not crappy."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "bs=64\n",
+ "sz=128"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen=None\n",
+ "gc.collect()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "loss_critic = AdaptiveLoss(nn.BCEWithLogitsLoss())"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "data_crit = get_crit_data([name_gen, 'test'], bs=bs, sz=sz)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "data_crit.show_batch(rows=3, ds_type=DatasetType.Train, imgsize=3)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_critic = colorize_crit_learner(data=data_crit, nf=256)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_critic.callback_fns.append(partial(LearnerTensorboardWriter, base_dir=TENSORBOARD_PATH, name='CriticPre'))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_critic.fit_one_cycle(6, 1e-3)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_critic.save(crit_name)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "bs=16\n",
+ "sz=192"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_critic.data=get_crit_data([name_gen, 'test'], bs=bs, sz=sz)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_critic.fit_one_cycle(4, 1e-4)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_critic.save(crit_name)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## GAN"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Now we'll combine those pretrained model in a GAN."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_crit=None\n",
+ "learn_gen=None\n",
+ "gc.collect()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "lr=2e-5\n",
+ "sz=192\n",
+ "bs=5"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "#placeholder- not actually used\n",
+ "data_crit = get_crit_data([name_gen, 'test'], bs=bs, sz=sz)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_crit = colorize_crit_learner(data=data_crit, nf=256).load(crit_name, with_opt=False)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen = gen_learner_wide(arch=models.resnet101, data=data_gen, gen_loss=FeatureLoss2(), nf_factor=nf_factor).load(gen_name, with_opt=False)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "switcher = partial(AdaptiveGANSwitcher, critic_thresh=0.65)\n",
+ "learn = GANLearner.from_learners(learn_gen, learn_crit, weights_gen=(1.0,1.5), show_img=False, switcher=switcher,\n",
+ " opt_func=partial(optim.Adam, betas=(0.,0.9)), wd=1e-3)\n",
+ "learn.callback_fns.append(partial(GANDiscriminativeLR, mult_lr=5.))\n",
+ "learn.callback_fns.append(partial(GANTensorboardWriter, base_dir=TENSORBOARD_PATH, name='GanLearner', visual_iters=100))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "for i in range(1,101):\n",
+ " learn.data = get_data(sz=sz, bs=bs, keep_pct=0.001)\n",
+ " learn_gen.freeze_to(-1)\n",
+ " learn.fit(1,lr)\n",
+ " save_all('_03_' + str(i))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "save_all('_01')"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Save Generated Images Again"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "bs=8\n",
+ "sz=192"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen = gen_learner_wide(arch=models.resnet101, data=data_gen, gen_loss=FeatureLoss2(), nf_factor=nf_factor).load('ColorizeNew73_gen192_05_7', with_opt=False)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "save_gen_images(gen_name)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### Train Critic Again"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "bs=16\n",
+ "sz=192"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen=None\n",
+ "gc.collect()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "loss_critic = AdaptiveLoss(nn.BCEWithLogitsLoss())"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "data_crit = get_crit_data([name_gen, 'test'], bs=bs, sz=sz)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "data_crit.show_batch(rows=3, ds_type=DatasetType.Train, imgsize=3)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_critic = colorize_crit_learner(data=data_crit, nf=256).load(crit_name + '5', with_opt=False)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_critic.callback_fns.append(partial(LearnerTensorboardWriter, base_dir=TENSORBOARD_PATH, name='CriticPre'))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_critic.fit_one_cycle(4, 1e-4)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_critic.save(crit_name + '6')"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "### GAN Again"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_crit=None\n",
+ "learn_gen=None\n",
+ "gc.collect()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "lr=1e-6\n",
+ "sz=192\n",
+ "bs=5"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "data_crit = get_crit_data([name_gen, 'test'], bs=bs, sz=sz)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_crit = colorize_crit_learner(data=data_crit, nf=256).load(crit_name + '6', with_opt=False)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "learn_gen = gen_learner_wide(arch=models.resnet101, data=data_gen, gen_loss=FeatureLoss2(), nf_factor=nf_factor).load('ColorizeNew73_gen192_05_7', with_opt=False)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "switcher = partial(AdaptiveGANSwitcher, critic_thresh=0.65)\n",
+ "learn = GANLearner.from_learners(learn_gen, learn_crit, weights_gen=(1.0,1.5), show_img=False, switcher=switcher,\n",
+ " opt_func=partial(optim.Adam, betas=(0.,0.9)), wd=1e-3)\n",
+ "learn.callback_fns.append(partial(GANDiscriminativeLR, mult_lr=5.))\n",
+ "learn.callback_fns.append(partial(GANTensorboardWriter, base_dir=TENSORBOARD_PATH, name='GanLearner', visual_iters=100))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "for i in range(1,101):\n",
+ " learn.data = get_data(sz=sz, bs=bs, keep_pct=0.001)\n",
+ " learn_gen.freeze_to(-1)\n",
+ " learn.fit(1,lr)\n",
+ " save_all('_06_' + str(i))"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## fin"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.7.0"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/ImageColorizer.ipynb b/ImageColorizer.ipynb
index e6e9b7a0..5afea3d2 100644
--- a/ImageColorizer.ipynb
+++ b/ImageColorizer.ipynb
@@ -42,13 +42,7 @@
"#It literally just is a number multiplied by 16 to get the square render resolution. \n",
"#Note that this doesn't affect the resolution of the final output- the output is the same resolution as the input.\n",
"#Example: render_factor=21 => color is rendered at 16x21 = 336x336 px. \n",
- "render_factor=17\n",
- "root_folder = Path('data/imagenet/ILSVRC/Data/CLS-LOC/bandw')\n",
- "weights_name = 'ColorizeNew68_gen192_01_5'\n",
- "nf_factor = 1.25\n",
- "\n",
- "#weights_name = 'ColorizeNew70_gen192_01_5'\n",
- "#nf_factor = 1.25"
+ "render_factor=48"
]
},
{
@@ -57,8 +51,10 @@
"metadata": {},
"outputs": [],
"source": [
- "vis = get_colorize_visualizer(root_folder=root_folder, weights_name=weights_name, nf_factor=nf_factor, render_factor=render_factor)\n",
- "#vis = get_colorize_visualizer(root_folder=root_folder, weights_name=weights_name, nf_factor=nf_factor, render_factor=render_factor)"
+ "vis = get_video_colorizer(root_folder=Path('data/imagenet/ILSVRC/Data/CLS-LOC/bandw'), weights_name='ColorizeNew76_gen192_01_28', render_factor=render_factor, nf_factor=1.5).vis\n",
+ "#vis = get_video_colorizer(root_folder=Path('data/imagenet/ILSVRC/Data/CLS-LOC/bandw'), weights_name='ColorizeNew72_gen192_05_18', render_factor=render_factor).vis\n",
+ "#vis = get_image_colorizer(render_factor=render_factor)\n",
+ "#vis = get_image_colorizer(arch=models.resnet101, root_folder=Path('data/imagenet/ILSVRC/Data/CLS-LOC/bandw'), weights_name='ColorizeNew73_gen192_06_80', render_factor=render_factor)"
]
},
{
@@ -1462,7 +1458,7 @@
"metadata": {},
"outputs": [],
"source": [
- "vis.plot_transformed_image(\"test_images/LivingRoom1920Sweeden.jpg\")"
+ "vis.plot_transformed_image(\"test_images/LivingRoom1920Sweden.jpg\")"
]
},
{
@@ -2407,7 +2403,7 @@
"metadata": {},
"outputs": [],
"source": [
- "vis.plot_transformed_image(\"test_images/WaterfordIreland1909.jpg\", figsize=(70,70))"
+ "vis.plot_transformed_image(\"test_images/WaterfordIreland1909.jpg\")"
]
},
{
diff --git a/README.md b/README.md
index 9b29821e..dab0786b 100644
--- a/README.md
+++ b/README.md
@@ -94,13 +94,13 @@ Seneca Native in 1908
This is a deep learning based model. More specifically, what I've done is combined the following approaches:
* **Self-Attention Generative Adversarial Network** (https://arxiv.org/abs/1805.08318). Except the generator is a **pretrained U-Net**, and I've just modified it to have the spectral normalization and self-attention. It's a pretty straightforward translation.
-* **Two Time-Scale Update Rule** (https://arxiv.org/abs/1706.08500). This is also very straightforward – it's just one to one generator/critic iterations and higher critic learning rate. This is modified to incorporate a "threshold" critic loss that makes sure that the critic is "caught up" before moving on to generator training. This is particularly useful for the GAN supertransfer learning method described next.
-* **GAN Supertransfer Learning** There's no paper here! And I just totally made up that catchy term. But it's the best way I can describe it. Basically what you do is you first train the generator in a conventional way by itself with just the feature loss. Then you generate images from that, and train the critic on distinguishing between those outputs and real images as a basic binary classifier. Finally, you train the generator and critic together in a GAN setting (starting right at the target size of 192px in this case). This training is super quick- only 1-10% of Imagenet dataset is iterated through, once! Yet during this very short amount of GAN training the generator not only gets the full realistic colorization capabilities that we used to get through days of progressively resized GAN training, but it also doesn't accrue any of the artifacts and other ugly baggage of GANs. As far as I know this is a new technique. And it's incredibly effective. It seems paper-worthy but I'll leave the paper to whoever's so inclined (not I!). This builds upon a technique developed in collaboration with Jeremy Howard and Sylvain Gugger (so fun!) for Fast.AI's Lesson 7 in version 3 of Practical Deep Learning for Coders part I. The particular lesson notebook can be found here: https://github.com/fastai/course-v3/blob/master/nbs/dl1/lesson7-superres-gan.ipynb
+* **Two Time-Scale Update Rule** (https://arxiv.org/abs/1706.08500). This is also very straightforward – it's just one to one generator/critic iterations and higher critic learning rate. This is modified to incorporate a "threshold" critic loss that makes sure that the critic is "caught up" before moving on to generator training. This is particularly useful for the GAN Supertransfer Learning method described next.
+* **GAN Supertransfer Learning**. There's no paper here! And I just totally made up that catchy term. But it's the best way I can describe it. Basically what you do is you first train the generator in a conventional way by itself with just the feature loss. Then you generate images from that, and train the critic on distinguishing between those outputs and real images as a basic binary classifier. Finally, you train the generator and critic together in a GAN setting (starting right at the target size of 192px in this case). This training is super quick- only 1-10% of Imagenet dataset is iterated through, once! Yet during this very short amount of GAN training the generator not only gets the full realistic colorization capabilities that we used to get through days of progressively resized GAN training, but it also doesn't accrue any of the artifacts and other ugly baggage of GANs. As far as I know this is a new technique. And it's incredibly effective. It seems paper-worthy but I'll leave the paper to whoever is so inclined (not I!). This builds upon a technique developed in collaboration with Jeremy Howard and Sylvain Gugger (so fun!) for Fast.AI's Lesson 7 in version 3 of Practical Deep Learning for Coders part I. The particular lesson notebook can be found here: https://github.com/fastai/course-v3/blob/master/nbs/dl1/lesson7-superres-gan.ipynb
* **Generator Loss** during GAN Supertransfer Learning is two parts: One is a basic Perceptual Loss (or Feature Loss) based on VGG16 – this just biases the generator model to replicate the input image. The second is the loss score from the critic. For the curious – Perceptual Loss isn't sufficient by itself to produce good results. It tends to just encourage a bunch of brown/green/blue – you know, cheating to the test, basically, which neural networks are really good at doing! Key thing to realize here is that GANs essentially are learning the loss function for you – which is really one big step closer to toward the ideal that we're shooting for in machine learning. And of course you generally get much better results when you get the machine to learn something you were previously hand coding. That's certainly the case here.
Of note: There's no longer any "Progressive Growing of GANs" type training going on here. It's just not needed in lieu of the superior results obtained by the GAN Supertransfer Learning technique described above.
-The beauty of this model is that it should be generally useful for all sorts of image modification, and it should do it quite well. What you're seeing above are the results of the colorization model, but that's just one component in a pipeline that I'm looking to develop here with the exact same model.
+The beauty of this model is that it should be generally useful for all sorts of image modification, and it should do it quite well. What you're seeing above are the results of the colorization model, but that's just one component in a pipeline that I'm looking to develop here with the exact same approach.
### This Project, Going Forward
@@ -115,7 +115,7 @@ The easiest way to get started is to go straight to the Colab notebooks:
Image [](https://colab.research.google.com/github/jantic/DeOldify/blob/master/ImageColorizerColab.ipynb) | Video [](https://colab.research.google.com/github/jantic/DeOldify/blob/master/VideoColorizerColab.ipynb)
-Special thanks to Matt Robinson and Maria Benevente for their image Colab notebook contributions, and Robert Bell for the video Colab notebook work!
+Special thanks to Matt Robinson and María Benavente for their image Colab notebook contributions, and Robert Bell for the video Colab notebook work!
-----------------------
@@ -123,8 +123,8 @@ Special thanks to Matt Robinson and Maria Benevente for their image Colab notebo
#### Hardware and Operating System Requirements
-* **(Training Only) BEEFY Graphics card**. I'd really like to have more memory than the 11 GB in my GeForce 1080TI (11GB). You'll have a tough time with less. The Unet and Critic are ridiculously large.
-* **(Colorization Alone) A decent graphics card**. Approximately 3GB+ memory video cards should be sufficient.
+* **(Training Only) BEEFY Graphics card**. I'd really like to have more memory than the 11 GB in my GeForce 1080TI (11GB). You'll have a tough time with less. The Generators and Critic are ridiculously large.
+* **(Colorization Alone) A decent graphics card**. Approximately 4GB+ memory video cards should be sufficient.
* **Linux (or maybe Windows 10)** I'm using Ubuntu 16.04, but nothing about this precludes Windows 10 support as far as I know. I just haven't tested it and am not going to make it a priority for now.
#### Easy Install
diff --git a/VideoColorizer.ipynb b/VideoColorizer.ipynb
index f19cd89e..e044fd2b 100644
--- a/VideoColorizer.ipynb
+++ b/VideoColorizer.ipynb
@@ -2,7 +2,7 @@
"cells": [
{
"cell_type": "code",
- "execution_count": null,
+ "execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
@@ -12,7 +12,7 @@
},
{
"cell_type": "code",
- "execution_count": null,
+ "execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
@@ -36,7 +36,7 @@
},
{
"cell_type": "code",
- "execution_count": null,
+ "execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
@@ -44,188 +44,110 @@
"#It literally just is a number multiplied by 16 to get the square render resolution. \n",
"#Note that this doesn't affect the resolution of the final output- the output is the same resolution as the input.\n",
"#Example: render_factor=21 => color is rendered at 16x21 = 336x336 px. \n",
- "render_factor=21\n",
- "root_folder = Path('data/imagenet/ILSVRC/Data/CLS-LOC/bandw')\n",
- "#weights_name = 'ColorizeNew50_gen192_10'\n",
- "weights_name = 'ColorizeNew68_gen192_01_5'\n",
- "nf_factor = 1.25\n",
+ "render_factor=43\n",
"\n",
- "workfolder = Path('./video')\n",
- "source_folder = workfolder/\"source\"\n",
- "bwframes_root = workfolder/\"bwframes\"\n",
- "colorframes_root = workfolder/\"colorframes\"\n",
- "result_folder = workfolder/\"result\"\n",
- "#Make source_url None to just read from source_path directly without modification\n",
- "source_url = 'https://twitter.com/silentmoviegifs/status/1092793719173115905'\n",
- "#source_url=None\n",
- "source_name = 'video8.mp4'\n",
- "source_path = source_folder/source_name\n",
- "bwframes_folder = bwframes_root/(source_path.stem)\n",
- "colorframes_folder = colorframes_root/(source_path.stem)\n",
- "result_path = result_folder/source_name"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "def progress(value, max=100):\n",
- " return HTML(\"\"\"\n",
- " \n",
- " \"\"\".format(value=value, max=max))"
+ "#Specify media_url. Many sources will work (YouTube, Imgur, Twitter, Reddit, etc). \n",
+ "#Complete list here: https://rg3.github.io/youtube-dl/supportedsites.html . \n",
+ "#NOTE: Make source_url None to just read from file at ./video/source/[file_name] directly without modification\n",
+ "#source_url = 'https://vimeo.com/87890004'\n",
+ "#source_url = 'https://www.youtube.com/watch?v=gZShc8oshtU'\n",
+ "#source_url = 'https://www.youtube.com/watch?v=fk6qiJjEEBo'\n",
+ "#source_url = 'https://twitter.com/silentmoviegifs/status/1088830101863759872'\n",
+ "source_url = None\n",
+ "file_name = 'video14.mp4'"
]
},
{
"cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "def get_fps():\n",
- " probe = ffmpeg.probe(str(source_path))\n",
- " stream_data = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)\n",
- " avg_frame_rate = stream_data['avg_frame_rate']\n",
- " print(avg_frame_rate)\n",
- " fps_num=avg_frame_rate.split(\"/\")[0]\n",
- " fps_den = avg_frame_rate.rsplit(\"/\")[1]\n",
- " return round(float(fps_num)/float(fps_den))"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "def purge_images(dir):\n",
- " for f in os.listdir(dir):\n",
- " if re.search('.*?\\.jpg', f):\n",
- " os.remove(os.path.join(dir, f))"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "## Download Video (optional via setting source_url)"
- ]
- },
- {
- "cell_type": "markdown",
+ "execution_count": 4,
"metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "/media/jason/Projects/Deep Learning/DeOldifyV2/DeOldify/fastai/data_block.py:426: UserWarning: Your training set is empty. If this is by design, pass `ignore_empty=True` to remove this warning.\n",
+ " warn(\"Your training set is empty. If this is by design, pass `ignore_empty=True` to remove this warning.\")\n",
+ "/media/jason/Projects/Deep Learning/DeOldifyV2/DeOldify/fastai/data_block.py:429: UserWarning: Your validation set is empty. If this is by design, use `no_split()`\n",
+ " or pass `ignore_empty=True` when labelling to remove this warning.\n",
+ " or pass `ignore_empty=True` when labelling to remove this warning.\"\"\")\n"
+ ]
+ }
+ ],
"source": [
- "##### Specify media_url. Many sources will work (YouTube, Imgur, Twitter, Reddit, etc). Complete list here: https://rg3.github.io/youtube-dl/supportedsites.html . The resulting file path can be used later."
+ "#colorizer = get_video_colorizer(render_factor=render_factor)\n",
+ "colorizer = get_video_colorizer2(render_factor=render_factor)"
]
},
{
"cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
+ "execution_count": 5,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ "
\n",
+ " \n",
+ "
\n",
+ " Interrupted\n",
+ "
\n",
+ " "
+ ],
+ "text/plain": [
+ ""
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "ename": "KeyboardInterrupt",
+ "evalue": "",
+ "output_type": "error",
+ "traceback": [
+ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+ "\u001b[0;31mKeyboardInterrupt\u001b[0m Traceback (most recent call last)",
+ "\u001b[0;32m\u001b[0m in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 2\u001b[0m \u001b[0mcolorizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcolorize_from_url\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msource_url\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfile_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 3\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 4\u001b[0;31m \u001b[0mcolorizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcolorize_from_file_name\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfile_name\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
+ "\u001b[0;32m/media/jason/Projects/Deep Learning/DeOldifyV2/DeOldify/fasterai/visualize.py\u001b[0m in \u001b[0;36mcolorize_from_file_name\u001b[0;34m(self, file_name)\u001b[0m\n\u001b[1;32m 128\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mcolorize_from_file_name\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfile_name\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 129\u001b[0m \u001b[0msource_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msource_folder\u001b[0m\u001b[0;34m/\u001b[0m\u001b[0mfile_name\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 130\u001b[0;31m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_colorize_from_path\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msource_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 131\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 132\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_colorize_from_path\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msource_path\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mPath\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m/media/jason/Projects/Deep Learning/DeOldifyV2/DeOldify/fasterai/visualize.py\u001b[0m in \u001b[0;36m_colorize_from_path\u001b[0;34m(self, source_path)\u001b[0m\n\u001b[1;32m 132\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_colorize_from_path\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msource_path\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mPath\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 133\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_extract_raw_frames\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msource_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 134\u001b[0;31m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_colorize_raw_frames\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msource_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 135\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_build_video\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msource_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 136\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m/media/jason/Projects/Deep Learning/DeOldifyV2/DeOldify/fasterai/visualize.py\u001b[0m in \u001b[0;36m_colorize_raw_frames\u001b[0;34m(self, source_path)\u001b[0m\n\u001b[1;32m 104\u001b[0m \u001b[0mimg_path\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbwframes_folder\u001b[0m\u001b[0;34m/\u001b[0m\u001b[0mimg\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 105\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0misfile\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimg_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 106\u001b[0;31m \u001b[0mcolor_image\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvis\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_transformed_image\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mimg_path\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 107\u001b[0m \u001b[0mcolor_image\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcolorframes_folder\u001b[0m\u001b[0;34m/\u001b[0m\u001b[0mimg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 108\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m/media/jason/Projects/Deep Learning/DeOldifyV2/DeOldify/fasterai/visualize.py\u001b[0m in \u001b[0;36mget_transformed_image\u001b[0;34m(self, path, render_factor)\u001b[0m\n\u001b[1;32m 39\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mget_transformed_image\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpath\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mPath\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrender_factor\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m->\u001b[0m\u001b[0mImage\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 40\u001b[0m \u001b[0morig_image\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_open_pil_image\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 41\u001b[0;31m \u001b[0mfiltered_image\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfilter\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfilter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0morig_image\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0morig_image\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrender_factor\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mrender_factor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 42\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mfiltered_image\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 43\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m/media/jason/Projects/Deep Learning/DeOldifyV2/DeOldify/fasterai/filters.py\u001b[0m in \u001b[0;36mfilter\u001b[0;34m(self, orig_image, filtered_image, render_factor)\u001b[0m\n\u001b[1;32m 101\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 102\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mfilter\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfilters\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 103\u001b[0;31m \u001b[0mfiltered_image\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfilter\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfilter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0morig_image\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfiltered_image\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrender_factor\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 104\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 105\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mfiltered_image\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m/media/jason/Projects/Deep Learning/DeOldifyV2/DeOldify/fasterai/filters.py\u001b[0m in \u001b[0;36mfilter\u001b[0;34m(self, orig_image, filtered_image, render_factor)\u001b[0m\n\u001b[1;32m 64\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mfilter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0morig_image\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mPilImage\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfiltered_image\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mPilImage\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrender_factor\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mint\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m->\u001b[0m\u001b[0mPilImage\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 65\u001b[0m \u001b[0mrender_sz\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mrender_factor\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrender_base\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 66\u001b[0;31m \u001b[0mmodel_image\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_model_process\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0morig\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfiltered_image\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msz\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mrender_sz\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 67\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 68\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmap_to_orig\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m/media/jason/Projects/Deep Learning/DeOldifyV2/DeOldify/fasterai/filters.py\u001b[0m in \u001b[0;36m_model_process\u001b[0;34m(self, orig, sz)\u001b[0m\n\u001b[1;32m 44\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnorm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdo_x\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 45\u001b[0m result = self.learn.pred_batch(ds_type=DatasetType.Valid, \n\u001b[0;32m---> 46\u001b[0;31m batch=(x[None].cuda(),y[None]), reconstruct=True)\n\u001b[0m\u001b[1;32m 47\u001b[0m \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 48\u001b[0m \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdenorm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdo_x\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m/media/jason/Projects/Deep Learning/DeOldifyV2/DeOldify/fastai/basic_train.py\u001b[0m in \u001b[0;36mpred_batch\u001b[0;34m(self, ds_type, batch, reconstruct)\u001b[0m\n\u001b[1;32m 315\u001b[0m \u001b[0mcb_handler\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mCallbackHandler\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcallbacks\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 316\u001b[0m \u001b[0mxb\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0myb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcb_handler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_batch_begin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mxb\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0myb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 317\u001b[0;31m \u001b[0mpreds\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mloss_batch\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0meval\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mxb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0myb\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcb_handler\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcb_handler\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 318\u001b[0m \u001b[0mres\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_loss_func2activ\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mloss_func\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpreds\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 319\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mreconstruct\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mres\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m/media/jason/Projects/Deep Learning/DeOldifyV2/DeOldify/fastai/basic_train.py\u001b[0m in \u001b[0;36mloss_batch\u001b[0;34m(model, xb, yb, loss_func, opt, cb_handler)\u001b[0m\n\u001b[1;32m 22\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mis_listy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mxb\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mxb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mxb\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 23\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mis_listy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0myb\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0myb\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0myb\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 24\u001b[0;31m \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mxb\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 25\u001b[0m \u001b[0mout\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcb_handler\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mon_loss_begin\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mout\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 26\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m~/anaconda3/envs/fastaiv1/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m 487\u001b[0m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 488\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 489\u001b[0;31m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 490\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 491\u001b[0m \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m/media/jason/Projects/Deep Learning/DeOldifyV2/DeOldify/fastai/layers.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m 131\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0ml\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlayers\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 132\u001b[0m \u001b[0mres\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0morig\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 133\u001b[0;31m \u001b[0mnres\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0ml\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mres\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 134\u001b[0m \u001b[0;31m# We have to remove res.orig to avoid hanging refs and therefore memory leaks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 135\u001b[0m \u001b[0mres\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0morig\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m~/anaconda3/envs/fastaiv1/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m 487\u001b[0m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 488\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 489\u001b[0;31m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 490\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 491\u001b[0m \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m~/anaconda3/envs/fastaiv1/lib/python3.7/site-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m 90\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 91\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_modules\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 92\u001b[0;31m \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 93\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 94\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m~/anaconda3/envs/fastaiv1/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m 487\u001b[0m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 488\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 489\u001b[0;31m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 490\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 491\u001b[0m \u001b[0mhook_result\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m~/anaconda3/envs/fastaiv1/lib/python3.7/site-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m 90\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 91\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_modules\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 92\u001b[0;31m \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 93\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 94\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m~/anaconda3/envs/fastaiv1/lib/python3.7/site-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m 483\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 484\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mhook\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_forward_pre_hooks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mvalues\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 485\u001b[0;31m \u001b[0mhook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 486\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_C\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_get_tracing_state\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 487\u001b[0m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_slow_forward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m~/anaconda3/envs/fastaiv1/lib/python3.7/site-packages/torch/nn/utils/spectral_norm.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, module, inputs)\u001b[0m\n\u001b[1;32m 98\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 99\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m__call__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 100\u001b[0;31m \u001b[0msetattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodule\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcompute_weight\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodule\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdo_power_iteration\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtraining\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 101\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 102\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_solve_v_and_rescale\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mweight_mat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mu\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtarget_sigma\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;32m~/anaconda3/envs/fastaiv1/lib/python3.7/site-packages/torch/nn/utils/spectral_norm.py\u001b[0m in \u001b[0;36mcompute_weight\u001b[0;34m(self, module, do_power_iteration)\u001b[0m\n\u001b[1;32m 84\u001b[0m \u001b[0mv\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mv\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mclone\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 85\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 86\u001b[0;31m \u001b[0msigma\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdot\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mu\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mweight_mat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mv\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 87\u001b[0m \u001b[0mweight\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mweight\u001b[0m \u001b[0;34m/\u001b[0m \u001b[0msigma\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 88\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mweight\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+ "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+ ]
+ }
+ ],
"source": [
"if source_url is not None:\n",
- " if source_path.exists(): source_path.unlink()\n",
- " youtubdl_command = 'youtube-dl \"' + source_url + '\" -o \"' + str(source_path) + '\"'\n",
- " print(youtubdl_command)\n",
- " print('\\n')\n",
- " output = Path(os.popen(youtubdl_command).read())\n",
- " print(str(output))"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "## Extract Raw Frames"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "bwframe_path_template = str(bwframes_folder/'%5d.jpg')\n",
- "bwframes_folder.mkdir(parents=True, exist_ok=True)\n",
- "purge_images(bwframes_folder)\n",
- "ffmpeg.input(str(source_path)).output(str(bwframe_path_template), format='image2', vcodec='mjpeg', qscale=0).run(capture_stdout=True)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "framecount = len(os.listdir(str(bwframes_folder)))"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "## DeOldify / Colorize"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "vis = get_colorize_visualizer(root_folder=root_folder, weights_name=weights_name, nf_factor=nf_factor, render_factor=render_factor)\n",
- "#vis = get_colorize_visualizer(render_factor=render_factor)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "prog = 0\n",
- "out = display(progress(0, 100), display_id=True)\n",
- "colorframes_folder.mkdir(parents=True, exist_ok=True)\n",
- "purge_images(colorframes_folder)\n",
- "\n",
- "for img in os.listdir(str(bwframes_folder)):\n",
- " img_path = bwframes_folder/img\n",
- " if os.path.isfile(str(img_path)):\n",
- " color_image = vis.get_transformed_image(str(img_path), render_factor)\n",
- " color_image.save(str(colorframes_folder/img))\n",
- " prog += 1\n",
- " out.update(progress(prog, framecount))"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "## Build Video"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "colorframes_path_template = str(colorframes_folder/'%5d.jpg')\n",
- "result_path.parent.mkdir(parents=True, exist_ok=True)\n",
- "\n",
- "if result_path.exists(): result_path.unlink()\n",
- "\n",
- "fps = get_fps()\n",
- "print(fps)\n",
- "ffmpeg.input(str(colorframes_path_template), format='image2', vcodec='mjpeg', framerate=str(fps)).output(str(result_path), crf=17, vcodec='libx264').run(capture_stdout=True)"
+ " colorizer.colorize_from_url(source_url, file_name)\n",
+ "else:\n",
+ " colorizer.colorize_from_file_name(file_name)"
]
},
{
diff --git a/fasterai/generators.py b/fasterai/generators.py
index 2a65c408..62b028dd 100644
--- a/fasterai/generators.py
+++ b/fasterai/generators.py
@@ -1,31 +1,31 @@
from fastai.vision import *
from fastai.vision.learner import cnn_config
-from .unet import CustomDynamicUnet, CustomDynamicUnet2
+from .unet import DynamicUnetWide, DynamicUnetDeep
from .loss import FeatureLoss
from .dataset import *
#Weights are implicitly read from ./models/ folder
-def colorize_gen_inference(root_folder:Path, weights_name:str, nf_factor:float)->Learner:
+def gen_inference_deep(root_folder:Path, weights_name:str, arch=models.resnet34, nf_factor:float=1.25)->Learner:
data = get_dummy_databunch()
- learn = colorize_gen_learner(data=data, gen_loss=F.l1_loss, nf_factor=nf_factor)
+ learn = gen_learner_deep(data=data, gen_loss=F.l1_loss, arch=arch, nf_factor=nf_factor)
learn.path = root_folder
learn.load(weights_name)
learn.model.eval()
return learn
-def colorize_gen_learner(data:ImageDataBunch, gen_loss=FeatureLoss(), arch=models.resnet34, nf_factor:float=1.0)->Learner:
- return custom_unet_learner(data, arch, wd=1e-3, blur=True, norm_type=NormType.Spectral,
+def gen_learner_deep(data:ImageDataBunch, gen_loss=FeatureLoss(), arch=models.resnet34, nf_factor:float=1.25)->Learner:
+ return unet_learner_deep(data, arch, wd=1e-3, blur=True, norm_type=NormType.Spectral,
self_attention=True, y_range=(-3.,3.), loss_func=gen_loss, nf_factor=nf_factor)
#The code below is meant to be merged into fastaiv1 ideally
-def custom_unet_learner(data:DataBunch, arch:Callable, pretrained:bool=True, blur_final:bool=True,
+def unet_learner_deep(data:DataBunch, arch:Callable, pretrained:bool=True, blur_final:bool=True,
norm_type:Optional[NormType]=NormType, split_on:Optional[SplitFuncOrIdxList]=None,
blur:bool=False, self_attention:bool=False, y_range:Optional[Tuple[float,float]]=None, last_cross:bool=True,
bottle:bool=False, nf_factor:float=1.0, **kwargs:Any)->Learner:
"Build Unet learner from `data` and `arch`."
meta = cnn_config(arch)
body = create_body(arch, pretrained)
- model = to_device(CustomDynamicUnet(body, n_classes=data.c, blur=blur, blur_final=blur_final,
+ model = to_device(DynamicUnetDeep(body, n_classes=data.c, blur=blur, blur_final=blur_final,
self_attention=self_attention, y_range=y_range, norm_type=norm_type, last_cross=last_cross,
bottle=bottle, nf_factor=nf_factor), data.device)
learn = Learner(data, model, **kwargs)
@@ -37,27 +37,27 @@ def custom_unet_learner(data:DataBunch, arch:Callable, pretrained:bool=True, blu
#-----------------------------
#Weights are implicitly read from ./models/ folder
-def colorize_gen_inference2(root_folder:Path, weights_name:str, nf_factor:int, arch=models.resnet34)->Learner:
+def gen_inference_wide(root_folder:Path, weights_name:str, nf_factor:int=2, arch=models.resnet34)->Learner:
data = get_dummy_databunch()
- learn = colorize_gen_learner2(data=data, gen_loss=F.l1_loss, nf_factor=nf_factor, arch=arch)
+ learn = gen_learner_wide(data=data, gen_loss=F.l1_loss, nf_factor=nf_factor, arch=arch)
learn.path = root_folder
learn.load(weights_name)
learn.model.eval()
return learn
-def colorize_gen_learner2(data:ImageDataBunch, gen_loss=FeatureLoss(), arch=models.resnet34, nf_factor:int=1)->Learner:
- return custom_unet_learner2(data, arch=arch, wd=1e-3, blur=True, norm_type=NormType.Spectral,
+def gen_learner_wide(data:ImageDataBunch, gen_loss=FeatureLoss(), arch=models.resnet34, nf_factor:int=2)->Learner:
+ return unet_learner_wide(data, arch=arch, wd=1e-3, blur=True, norm_type=NormType.Spectral,
self_attention=True, y_range=(-3.,3.), loss_func=gen_loss, nf_factor=nf_factor)
#The code below is meant to be merged into fastaiv1 ideally
-def custom_unet_learner2(data:DataBunch, arch:Callable, pretrained:bool=True, blur_final:bool=True,
+def unet_learner_wide(data:DataBunch, arch:Callable, pretrained:bool=True, blur_final:bool=True,
norm_type:Optional[NormType]=NormType, split_on:Optional[SplitFuncOrIdxList]=None,
blur:bool=False, self_attention:bool=False, y_range:Optional[Tuple[float,float]]=None, last_cross:bool=True,
bottle:bool=False, nf_factor:int=1, **kwargs:Any)->Learner:
"Build Unet learner from `data` and `arch`."
meta = cnn_config(arch)
body = create_body(arch, pretrained)
- model = to_device(CustomDynamicUnet2(body, n_classes=data.c, blur=blur, blur_final=blur_final,
+ model = to_device(DynamicUnetWide(body, n_classes=data.c, blur=blur, blur_final=blur_final,
self_attention=self_attention, y_range=y_range, norm_type=norm_type, last_cross=last_cross,
bottle=bottle, nf_factor=nf_factor), data.device)
learn = Learner(data, model, **kwargs)
diff --git a/fasterai/unet.py b/fasterai/unet.py
index e8bc0660..1ba75c72 100644
--- a/fasterai/unet.py
+++ b/fasterai/unet.py
@@ -7,7 +7,7 @@
#The code below is meant to be merged into fastaiv1 ideally
-__all__ = ['CustomDynamicUnet', 'CustomUnetBlock', 'CustomPixelShuffle_ICNR']
+__all__ = ['DynamicUnetDeep', 'DynamicUnetWide']
def _get_sfs_idxs(sizes:Sizes) -> List[int]:
"Get the indexes of the layers where the size of the activation changes."
@@ -35,7 +35,7 @@ def forward(self,x):
x = self.shuf(self.relu(self.conv(x)))
return self.blur(self.pad(x)) if self.blur else x
-class CustomUnetBlock(nn.Module):
+class UnetBlockDeep(nn.Module):
"A quasi-UNet block, using `PixelShuffle_ICNR upsampling`."
def __init__(self, up_in_c:int, x_in_c:int, hook:Hook, final_div:bool=True, blur:bool=False, leaky:float=None,
self_attention:bool=False, nf_factor:float=1.0, **kwargs):
@@ -59,7 +59,7 @@ def forward(self, up_in:Tensor) -> Tensor:
return self.conv2(self.conv1(cat_x))
-class CustomDynamicUnet(SequentialEx):
+class DynamicUnetDeep(SequentialEx):
"Create a U-Net from a given architecture."
def __init__(self, encoder:nn.Module, n_classes:int, blur:bool=False, blur_final=True, self_attention:bool=False,
y_range:Optional[Tuple[float,float]]=None, last_cross:bool=True, bottle:bool=False,
@@ -82,7 +82,7 @@ def __init__(self, encoder:nn.Module, n_classes:int, blur:bool=False, blur_final
up_in_c, x_in_c = int(x.shape[1]), int(sfs_szs[idx][1])
do_blur = blur and (not_final or blur_final)
sa = self_attention and (i==len(sfs_idxs)-3)
- unet_block = CustomUnetBlock(up_in_c, x_in_c, self.sfs[i], final_div=not_final, blur=blur, self_attention=sa,
+ unet_block = UnetBlockDeep(up_in_c, x_in_c, self.sfs[i], final_div=not_final, blur=blur, self_attention=sa,
norm_type=norm_type, extra_bn=extra_bn, nf_factor=nf_factor, **kwargs).eval()
layers.append(unet_block)
x = unet_block(x)
@@ -104,7 +104,7 @@ def __del__(self):
#------------------------------------------------------
-class CustomUnetBlock2(nn.Module):
+class UnetBlockWide(nn.Module):
"A quasi-UNet block, using `PixelShuffle_ICNR upsampling`."
def __init__(self, up_in_c:int, x_in_c:int, n_out:int, hook:Hook, final_div:bool=True, blur:bool=False, leaky:float=None,
self_attention:bool=False, **kwargs):
@@ -127,7 +127,7 @@ def forward(self, up_in:Tensor) -> Tensor:
return self.conv(cat_x)
-class CustomDynamicUnet2(SequentialEx):
+class DynamicUnetWide(SequentialEx):
"Create a U-Net from a given architecture."
def __init__(self, encoder:nn.Module, n_classes:int, blur:bool=False, blur_final=True, self_attention:bool=False,
y_range:Optional[Tuple[float,float]]=None, last_cross:bool=True, bottle:bool=False,
@@ -155,7 +155,7 @@ def __init__(self, encoder:nn.Module, n_classes:int, blur:bool=False, blur_final
n_out = nf if not_final else nf//2
- unet_block = CustomUnetBlock2(up_in_c, x_in_c, n_out, self.sfs[i], final_div=not_final, blur=blur, self_attention=sa,
+ unet_block = UnetBlockWide(up_in_c, x_in_c, n_out, self.sfs[i], final_div=not_final, blur=blur, self_attention=sa,
norm_type=norm_type, extra_bn=extra_bn, **kwargs).eval()
layers.append(unet_block)
x = unet_block(x)
diff --git a/fasterai/visualize.py b/fasterai/visualize.py
index 3bda6004..3def8aab 100644
--- a/fasterai/visualize.py
+++ b/fasterai/visualize.py
@@ -4,11 +4,13 @@
from matplotlib.figure import Figure
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from .filters import IFilter, MasterFilter, ColorizerFilter
-from .generators import colorize_gen_inference, colorize_gen_inference2
+from .generators import gen_inference_deep, gen_inference_wide
from IPython.display import display
from tensorboardX import SummaryWriter
from scipy import misc
from PIL import Image
+import ffmpeg
+import youtube_dl
class ModelImageVisualizer():
@@ -51,17 +53,106 @@ def _get_num_rows_columns(self, num_images:int, max_columns:int)->(int,int):
rows = rows if rows * columns == num_images else rows + 1
return rows, columns
+class VideoColorizer():
+ def __init__(self, vis:ModelImageVisualizer):
+ self.vis=vis
+ workfolder = Path('./video')
+ self.source_folder = workfolder/"source"
+ self.bwframes_root = workfolder/"bwframes"
+ self.audio_root = workfolder/"audio"
+ self.colorframes_root = workfolder/"colorframes"
+ self.result_folder = workfolder/"result"
+
+ def _purge_images(self, dir):
+ for f in os.listdir(dir):
+ if re.search('.*?\.jpg', f):
+ os.remove(os.path.join(dir, f))
+
+ def _get_fps(self, source_path: Path)->float:
+ probe = ffmpeg.probe(str(source_path))
+ stream_data = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
+ avg_frame_rate = stream_data['avg_frame_rate']
+ fps_num=avg_frame_rate.split("/")[0]
+ fps_den = avg_frame_rate.rsplit("/")[1]
+ return round(float(fps_num)/float(fps_den))
+
+ def _download_video_from_url(self, source_url, source_path:Path):
+ if source_path.exists(): source_path.unlink()
+
+ ydl_opts = {
+ 'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/mp4',
+ 'outtmpl': str(source_path)
+ }
+ with youtube_dl.YoutubeDL(ydl_opts) as ydl:
+ ydl.download([source_url])
+
+ def _extract_raw_frames(self, source_path:Path):
+ bwframes_folder = self.bwframes_root/(source_path.stem)
+ bwframe_path_template = str(bwframes_folder/'%5d.jpg')
+ bwframes_folder.mkdir(parents=True, exist_ok=True)
+ self._purge_images(bwframes_folder)
+ ffmpeg.input(str(source_path)).output(str(bwframe_path_template), format='image2', vcodec='mjpeg', qscale=0).run(capture_stdout=True)
+
+
+ def _colorize_raw_frames(self, source_path:Path):
+ colorframes_folder = self.colorframes_root/(source_path.stem)
+ colorframes_folder.mkdir(parents=True, exist_ok=True)
+ self._purge_images(colorframes_folder)
+ bwframes_folder = self.bwframes_root/(source_path.stem)
+
+ for img in progress_bar(os.listdir(str(bwframes_folder))):
+ img_path = bwframes_folder/img
+ if os.path.isfile(str(img_path)):
+ color_image = self.vis.get_transformed_image(str(img_path))
+ color_image.save(str(colorframes_folder/img))
+
+ def _build_video(self, source_path:Path):
+ result_path = self.result_folder/source_path.name
+ colorframes_folder = self.colorframes_root/(source_path.stem)
+ colorframes_path_template = str(colorframes_folder/'%5d.jpg')
+ result_path.parent.mkdir(parents=True, exist_ok=True)
+ if result_path.exists(): result_path.unlink()
+ fps = self._get_fps(source_path)
+
+ ffmpeg.input(str(colorframes_path_template), format='image2', vcodec='mjpeg', framerate=str(fps)) \
+ .output(str(result_path), crf=17, vcodec='libx264') \
+ .run(capture_stdout=True)
+
+ print('Video created here: ' + str(result_path))
+
+ def colorize_from_url(self, source_url, file_name:str):
+ source_path = self.source_folder/file_name
+ self._download_video_from_url(source_url, source_path)
+ self._colorize_from_path(source_path)
+
+ def colorize_from_file_name(self, file_name:str):
+ source_path = self.source_folder/file_name
+ self._colorize_from_path(source_path)
+
+ def _colorize_from_path(self, source_path:Path):
+ self._extract_raw_frames(source_path)
+ self._colorize_raw_frames(source_path)
+ self._build_video(source_path)
+
+
+def get_video_colorizer2(root_folder:Path=Path('./'), weights_name:str='ColorizeVideos_gen2',
+ results_dir = 'result_images', render_factor:int=36)->VideoColorizer:
+ learn = gen_inference_wide(root_folder=root_folder, weights_name=weights_name, arch=models.resnet101)
+ filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
+ vis = ModelImageVisualizer(filtr, results_dir=results_dir)
+ return VideoColorizer(vis)
+
-def get_colorize_visualizer(root_folder:Path=Path('./'), weights_name:str='colorize_gen',
- results_dir = 'result_images', nf_factor:float=1.25, render_factor:int=21)->ModelImageVisualizer:
- learn = colorize_gen_inference(root_folder=root_folder, weights_name=weights_name, nf_factor=nf_factor)
+def get_video_colorizer(root_folder:Path=Path('./'), weights_name:str='ColorizeVideos_gen',
+ results_dir = 'result_images', render_factor:int=21, nf_factor:float=1.25)->VideoColorizer:
+ learn = gen_inference_deep(root_folder=root_folder, weights_name=weights_name, nf_factor=nf_factor)
filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
vis = ModelImageVisualizer(filtr, results_dir=results_dir)
- return vis
+ return VideoColorizer(vis)
-def get_colorize_visualizer2(root_folder:Path=Path('./'), weights_name:str='colorize_gen',
- results_dir = 'result_images', nf_factor:int=1, render_factor:int=21, arch=models.resnet34)->ModelImageVisualizer:
- learn = colorize_gen_inference2(root_folder=root_folder, weights_name=weights_name, nf_factor=nf_factor, arch=arch)
+def get_image_colorizer(root_folder:Path=Path('./'), weights_name:str='ColorizeImages_gen',
+ results_dir = 'result_images', render_factor:int=21, arch=models.resnet34)->ModelImageVisualizer:
+ learn = gen_inference_wide(root_folder=root_folder, weights_name=weights_name, arch=arch)
filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
vis = ModelImageVisualizer(filtr, results_dir=results_dir)
return vis
@@ -69,4 +160,3 @@ def get_colorize_visualizer2(root_folder:Path=Path('./'), weights_name:str='colo
-