Upload 16 files

src/app_utils.py

@@ -0,0 +1,131 @@
+import os
+import requests
+import random
+import _thread as thread
+from uuid import uuid4
+import urllib
+
+import numpy as np
+import skimage
+from skimage.filters import gaussian
+from PIL import Image
+
+def compress_image(image, path_original):
+    size = 1920, 1080
+    width = 1920
+    height = 1080
+
+    name = os.path.basename(path_original).split('.')
+    first_name = os.path.join(os.path.dirname(path_original), name[0] + '.jpg')
+
+    if image.size[0] > width and image.size[1] > height:
+        image.thumbnail(size, Image.ANTIALIAS)
+        image.save(first_name, quality=85)
+    elif image.size[0] > width:
+        wpercent = (width/float(image.size[0]))
+        height = int((float(image.size[1])*float(wpercent)))
+        image = image.resize((width,height), Image.ANTIALIAS)
+        image.save(first_name,quality=85)
+    elif image.size[1] > height:
+        wpercent = (height/float(image.size[1]))
+        width = int((float(image.size[0])*float(wpercent)))
+        image = image.resize((width,height), Image.ANTIALIAS)
+        image.save(first_name, quality=85)
+    else:
+        image.save(first_name, quality=85)
+
+
+def convertToJPG(path_original):
+    img = Image.open(path_original)
+    name = os.path.basename(path_original).split('.')
+    first_name = os.path.join(os.path.dirname(path_original), name[0] + '.jpg')
+
+    if img.format == "JPEG":
+        image = img.convert('RGB')
+        compress_image(image, path_original)
+        img.close()
+
+    elif img.format == "GIF":
+        i = img.convert("RGBA")
+        bg = Image.new("RGBA", i.size)
+        image = Image.composite(i, bg, i)
+        compress_image(image, path_original)
+        img.close()
+
+    elif img.format == "PNG":
+        try:
+            image = Image.new("RGB", img.size, (255,255,255))
+            image.paste(img,img)
+            compress_image(image, path_original)
+        except ValueError:
+            image = img.convert('RGB')
+            compress_image(image, path_original)
+        
+        img.close()
+
+    elif img.format == "BMP":
+        image = img.convert('RGB')
+        compress_image(image, path_original)
+        img.close()
+
+
+
+def blur(image, x0, x1, y0, y1, sigma=1, multichannel=True):
+    y0, y1 = min(y0, y1), max(y0, y1)
+    x0, x1 = min(x0, x1), max(x0, x1)
+    im = image.copy()
+    sub_im = im[y0:y1,x0:x1].copy()
+    blur_sub_im = gaussian(sub_im, sigma=sigma, multichannel=multichannel)
+    blur_sub_im = np.round(255 * blur_sub_im)
+    im[y0:y1,x0:x1] = blur_sub_im
+    return im
+
+
+
+def download(url, filename):
+    data = requests.get(url).content
+    with open(filename, 'wb') as handler:
+        handler.write(data)
+
+    return filename
+
+
+def generate_random_filename(upload_directory, extension):
+    filename = str(uuid4())
+    filename = os.path.join(upload_directory, filename + "." + extension)
+    return filename
+
+
+def clean_me(filename):
+    if os.path.exists(filename):
+        os.remove(filename)
+
+
+def clean_all(files):
+    for me in files:
+        clean_me(me)
+
+
+def create_directory(path):
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+
+
+def get_model_bin(url, output_path):
+    # print('Getting model dir: ', output_path)
+    if not os.path.exists(output_path):
+        create_directory(output_path)
+
+        urllib.request.urlretrieve(url, output_path)
+
+        # cmd = "wget -O %s %s" % (output_path, url)
+        # print(cmd)
+        # os.system(cmd)
+
+    return output_path
+
+
+#model_list = [(url, output_path), (url, output_path)]
+def get_multi_model_bin(model_list):
+    for m in model_list:
+        thread.start_new_thread(get_model_bin, m)
+

src/deoldify/__init__.py

@@ -0,0 +1,3 @@
+from src.deoldify._device import _Device
+
+device = _Device()

src/deoldify/_device.py

@@ -0,0 +1,30 @@
+import os
+from enum import Enum
+from .device_id import DeviceId
+
+#NOTE:  This must be called first before any torch imports in order to work properly!
+
+class DeviceException(Exception):
+    pass
+
+class _Device:
+    def __init__(self):
+        self.set(DeviceId.CPU)
+
+    def is_gpu(self):
+        ''' Returns `True` if the current device is GPU, `False` otherwise. '''
+        return self.current() is not DeviceId.CPU
+  
+    def current(self):
+        return self._current_device
+
+    def set(self, device:DeviceId):     
+        if device == DeviceId.CPU:
+            os.environ['CUDA_VISIBLE_DEVICES']=''
+        else:
+            os.environ['CUDA_VISIBLE_DEVICES']=str(device.value)
+            import torch
+            torch.backends.cudnn.benchmark=False
+        
+        self._current_device = device    
+        return device

src/deoldify/augs.py

@@ -0,0 +1,29 @@
+import random
+
+from fastai.vision.image import TfmPixel
+
+# Contributed by Rani Horev. Thank you!
+def _noisify(
+    x, pct_pixels_min: float = 0.001, pct_pixels_max: float = 0.4, noise_range: int = 30
+):
+    if noise_range > 255 or noise_range < 0:
+        raise Exception("noise_range must be between 0 and 255, inclusively.")
+
+    h, w = x.shape[1:]
+    img_size = h * w
+    mult = 10000.0
+    pct_pixels = (
+        random.randrange(int(pct_pixels_min * mult), int(pct_pixels_max * mult)) / mult
+    )
+    noise_count = int(img_size * pct_pixels)
+
+    for ii in range(noise_count):
+        yy = random.randrange(h)
+        xx = random.randrange(w)
+        noise = random.randrange(-noise_range, noise_range) / 255.0
+        x[:, yy, xx].add_(noise)
+
+    return x
+
+
+noisify = TfmPixel(_noisify)

src/deoldify/critics.py

@@ -0,0 +1,44 @@
+from fastai.core import *
+from fastai.torch_core import *
+from fastai.vision import *
+from fastai.vision.gan import AdaptiveLoss, accuracy_thresh_expand
+
+_conv_args = dict(leaky=0.2, norm_type=NormType.Spectral)
+
+
+def _conv(ni: int, nf: int, ks: int = 3, stride: int = 1, **kwargs):
+    return conv_layer(ni, nf, ks=ks, stride=stride, **_conv_args, **kwargs)
+
+
+def custom_gan_critic(
+    n_channels: int = 3, nf: int = 256, n_blocks: int = 3, p: int = 0.15
+):
+    "Critic to train a `GAN`."
+    layers = [_conv(n_channels, nf, ks=4, stride=2), nn.Dropout2d(p / 2)]
+    for i in range(n_blocks):
+        layers += [
+            _conv(nf, nf, ks=3, stride=1),
+            nn.Dropout2d(p),
+            _conv(nf, nf * 2, ks=4, stride=2, self_attention=(i == 0)),
+        ]
+        nf *= 2
+    layers += [
+        _conv(nf, nf, ks=3, stride=1),
+        _conv(nf, 1, ks=4, bias=False, padding=0, use_activ=False),
+        Flatten(),
+    ]
+    return nn.Sequential(*layers)
+
+
+def colorize_crit_learner(
+    data: ImageDataBunch,
+    loss_critic=AdaptiveLoss(nn.BCEWithLogitsLoss()),
+    nf: int = 256,
+) -> Learner:
+    return Learner(
+        data,
+        custom_gan_critic(nf=nf),
+        metrics=accuracy_thresh_expand,
+        loss_func=loss_critic,
+        wd=1e-3,
+    )

src/deoldify/dataset.py

@@ -0,0 +1,48 @@
+import fastai
+from fastai import *
+from fastai.core import *
+from fastai.vision.transform import get_transforms
+from fastai.vision.data import ImageImageList, ImageDataBunch, imagenet_stats
+from .augs import noisify
+
+
+def get_colorize_data(
+    sz: int,
+    bs: int,
+    crappy_path: Path,
+    good_path: Path,
+    random_seed: int = None,
+    keep_pct: float = 1.0,
+    num_workers: int = 8,
+    stats: tuple = imagenet_stats,
+    xtra_tfms=[],
+) -> ImageDataBunch:
+    
+    src = (
+        ImageImageList.from_folder(crappy_path, convert_mode='RGB')
+        .use_partial_data(sample_pct=keep_pct, seed=random_seed)
+        .split_by_rand_pct(0.1, seed=random_seed)
+    )
+
+    data = (
+        src.label_from_func(lambda x: good_path / x.relative_to(crappy_path))
+        .transform(
+            get_transforms(
+                max_zoom=1.2, max_lighting=0.5, max_warp=0.25, xtra_tfms=xtra_tfms
+            ),
+            size=sz,
+            tfm_y=True,
+        )
+        .databunch(bs=bs, num_workers=num_workers, no_check=True)
+        .normalize(stats, do_y=True)
+    )
+
+    data.c = 3
+    return data
+
+
+def get_dummy_databunch() -> ImageDataBunch:
+    path = Path('./assets/dummy/')
+    return get_colorize_data(
+        sz=1, bs=1, crappy_path=path, good_path=path, keep_pct=0.001
+    )

src/deoldify/device_id.py

@@ -0,0 +1,12 @@
+from enum import IntEnum
+
+class DeviceId(IntEnum):
+    GPU0 = 0,
+    GPU1 = 1,
+    GPU2 = 2,
+    GPU3 = 3,
+    GPU4 = 4,
+    GPU5 = 5,
+    GPU6 = 6,
+    GPU7 = 7,
+    CPU = 99

src/deoldify/filters.py

@@ -0,0 +1,120 @@
+from numpy import ndarray
+from abc import ABC, abstractmethod
+from .critics import colorize_crit_learner
+from fastai.core import *
+from fastai.vision import *
+from fastai.vision.image import *
+from fastai.vision.data import *
+from fastai import *
+import math
+from scipy import misc
+import cv2
+from PIL import Image as PilImage
+
+
+class IFilter(ABC):
+    @abstractmethod
+    def filter(
+        self, orig_image: PilImage, filtered_image: PilImage, render_factor: int
+    ) -> PilImage:
+        pass
+
+
+class BaseFilter(IFilter):
+    def __init__(self, learn: Learner, stats: tuple = imagenet_stats):
+        super().__init__()
+        self.learn = learn
+        self.device = next(self.learn.model.parameters()).device
+        self.norm, self.denorm = normalize_funcs(*stats)
+
+    def _transform(self, image: PilImage) -> PilImage:
+        return image
+
+    def _scale_to_square(self, orig: PilImage, targ: int) -> PilImage:
+        # a simple stretch to fit a square really makes a big difference in rendering quality/consistency.
+        # I've tried padding to the square as well (reflect, symetric, constant, etc).  Not as good!
+        targ_sz = (targ, targ)
+        return orig.resize(targ_sz, resample=PIL.Image.BILINEAR)
+
+    def _get_model_ready_image(self, orig: PilImage, sz: int) -> PilImage:
+        result = self._scale_to_square(orig, sz)
+        result = self._transform(result)
+        return result
+
+    def _model_process(self, orig: PilImage, sz: int) -> PilImage:
+        model_image = self._get_model_ready_image(orig, sz)
+        x = pil2tensor(model_image, np.float32)
+        x = x.to(self.device)
+        x.div_(255)
+        x, y = self.norm((x, x), do_x=True)
+        
+        try:
+            result = self.learn.pred_batch(
+                ds_type=DatasetType.Valid, batch=(x[None], y[None]), reconstruct=True
+            )
+        except RuntimeError as rerr:
+            if 'memory' not in str(rerr):
+                raise rerr
+            print('Warning: render_factor was set too high, and out of memory error resulted. Returning original image.')
+            return model_image
+            
+        out = result[0]
+        out = self.denorm(out.px, do_x=False)
+        out = image2np(out * 255).astype(np.uint8)
+        return PilImage.fromarray(out)
+
+    def _unsquare(self, image: PilImage, orig: PilImage) -> PilImage:
+        targ_sz = orig.size
+        image = image.resize(targ_sz, resample=PIL.Image.BILINEAR)
+        return image
+
+
+class ColorizerFilter(BaseFilter):
+    def __init__(self, learn: Learner, stats: tuple = imagenet_stats):
+        super().__init__(learn=learn, stats=stats)
+        self.render_base = 16
+
+    def filter(
+        self, orig_image: PilImage, filtered_image: PilImage, render_factor: int, post_process: bool = True) -> PilImage:
+        render_sz = render_factor * self.render_base
+        model_image = self._model_process(orig=filtered_image, sz=render_sz)
+        raw_color = self._unsquare(model_image, orig_image)
+
+        if post_process:
+            return self._post_process(raw_color, orig_image)
+        else:
+            return raw_color
+
+    def _transform(self, image: PilImage) -> PilImage:
+        return image.convert('LA').convert('RGB')
+
+    # This takes advantage of the fact that human eyes are much less sensitive to
+    # imperfections in chrominance compared to luminance.  This means we can
+    # save a lot on memory and processing in the model, yet get a great high
+    # resolution result at the end.  This is primarily intended just for
+    # inference
+    def _post_process(self, raw_color: PilImage, orig: PilImage) -> PilImage:
+        color_np = np.asarray(raw_color)
+        orig_np = np.asarray(orig)
+        color_yuv = cv2.cvtColor(color_np, cv2.COLOR_BGR2YUV)
+        # do a black and white transform first to get better luminance values
+        orig_yuv = cv2.cvtColor(orig_np, cv2.COLOR_BGR2YUV)
+        hires = np.copy(orig_yuv)
+        hires[:, :, 1:3] = color_yuv[:, :, 1:3]
+        final = cv2.cvtColor(hires, cv2.COLOR_YUV2BGR)
+        final = PilImage.fromarray(final)
+        return final
+
+
+class MasterFilter(BaseFilter):
+    def __init__(self, filters: [IFilter], render_factor: int):
+        self.filters = filters
+        self.render_factor = render_factor
+
+    def filter(
+        self, orig_image: PilImage, filtered_image: PilImage, render_factor: int = None, post_process: bool = True) -> PilImage:
+        render_factor = self.render_factor if render_factor is None else render_factor
+        for filter in self.filters:
+            filtered_image = filter.filter(orig_image, filtered_image, render_factor, post_process)
+
+        return filtered_image

src/deoldify/generators.py

@@ -0,0 +1,151 @@
+from fastai.vision import *
+from fastai.vision.learner import cnn_config
+from .unet import DynamicUnetWide, DynamicUnetDeep
+from .loss import FeatureLoss
+from .dataset import *
+
+# Weights are implicitly read from ./models/ folder
+def gen_inference_wide(
+    root_folder: Path, weights_name: str, nf_factor: int = 2, arch=models.resnet101) -> Learner:
+    data = get_dummy_databunch()
+    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 gen_learner_wide(
+    data: ImageDataBunch, gen_loss, arch=models.resnet101, 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.0, 3.0),
+        loss_func=gen_loss,
+        nf_factor=nf_factor,
+    )
+
+
+# The code below is meant to be merged into fastaiv1 ideally
+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(
+        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)
+    learn.split(ifnone(split_on, meta['split']))
+    if pretrained:
+        learn.freeze()
+    apply_init(model[2], nn.init.kaiming_normal_)
+    return learn
+
+
+# ----------------------------------------------------------------------
+
+# Weights are implicitly read from ./models/ folder
+def gen_inference_deep(
+    root_folder: Path, weights_name: str, arch=models.resnet34, nf_factor: float = 1.5) -> Learner:
+    data = get_dummy_databunch()
+    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 gen_learner_deep(
+    data: ImageDataBunch, gen_loss, arch=models.resnet34, nf_factor: float = 1.5
+) -> Learner:
+    return unet_learner_deep(
+        data,
+        arch,
+        wd=1e-3,
+        blur=True,
+        norm_type=NormType.Spectral,
+        self_attention=True,
+        y_range=(-3.0, 3.0),
+        loss_func=gen_loss,
+        nf_factor=nf_factor,
+    )
+
+
+# The code below is meant to be merged into fastaiv1 ideally
+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.5,
+    **kwargs: Any
+) -> Learner:
+    "Build Unet learner from `data` and `arch`."
+    meta = cnn_config(arch)
+    body = create_body(arch, pretrained)
+    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)
+    learn.split(ifnone(split_on, meta['split']))
+    if pretrained:
+        learn.freeze()
+    apply_init(model[2], nn.init.kaiming_normal_)
+    return learn
+
+
+# -----------------------------

src/deoldify/layers.py

@@ -0,0 +1,48 @@
+from fastai.layers import *
+from fastai.torch_core import *
+from torch.nn.parameter import Parameter
+from torch.autograd import Variable
+
+
+# The code below is meant to be merged into fastaiv1 ideally
+
+
+def custom_conv_layer(
+    ni: int,
+    nf: int,
+    ks: int = 3,
+    stride: int = 1,
+    padding: int = None,
+    bias: bool = None,
+    is_1d: bool = False,
+    norm_type: Optional[NormType] = NormType.Batch,
+    use_activ: bool = True,
+    leaky: float = None,
+    transpose: bool = False,
+    init: Callable = nn.init.kaiming_normal_,
+    self_attention: bool = False,
+    extra_bn: bool = False,
+):
+    "Create a sequence of convolutional (`ni` to `nf`), ReLU (if `use_activ`) and batchnorm (if `bn`) layers."
+    if padding is None:
+        padding = (ks - 1) // 2 if not transpose else 0
+    bn = norm_type in (NormType.Batch, NormType.BatchZero) or extra_bn == True
+    if bias is None:
+        bias = not bn
+    conv_func = nn.ConvTranspose2d if transpose else nn.Conv1d if is_1d else nn.Conv2d
+    conv = init_default(
+        conv_func(ni, nf, kernel_size=ks, bias=bias, stride=stride, padding=padding),
+        init,
+    )
+    if norm_type == NormType.Weight:
+        conv = weight_norm(conv)
+    elif norm_type == NormType.Spectral:
+        conv = spectral_norm(conv)
+    layers = [conv]
+    if use_activ:
+        layers.append(relu(True, leaky=leaky))
+    if bn:
+        layers.append((nn.BatchNorm1d if is_1d else nn.BatchNorm2d)(nf))
+    if self_attention:
+        layers.append(SelfAttention(nf))
+    return nn.Sequential(*layers)

src/deoldify/loss.py

@@ -0,0 +1,136 @@
+from fastai import *
+from fastai.core import *
+from fastai.torch_core import *
+from fastai.callbacks import hook_outputs
+import torchvision.models as models
+
+
+class FeatureLoss(nn.Module):
+    def __init__(self, layer_wgts=[20, 70, 10]):
+        super().__init__()
+
+        self.m_feat = models.vgg16_bn(True).features.cuda().eval()
+        requires_grad(self.m_feat, False)
+        blocks = [
+            i - 1
+            for i, o in enumerate(children(self.m_feat))
+            if isinstance(o, nn.MaxPool2d)
+        ]
+        layer_ids = blocks[2:5]
+        self.loss_features = [self.m_feat[i] for i in layer_ids]
+        self.hooks = hook_outputs(self.loss_features, detach=False)
+        self.wgts = layer_wgts
+        self.metric_names = ['pixel'] + [f'feat_{i}' for i in range(len(layer_ids))]
+        self.base_loss = F.l1_loss
+
+    def _make_features(self, x, clone=False):
+        self.m_feat(x)
+        return [(o.clone() if clone else o) for o in self.hooks.stored]
+
+    def forward(self, input, target):
+        out_feat = self._make_features(target, clone=True)
+        in_feat = self._make_features(input)
+        self.feat_losses = [self.base_loss(input, target)]
+        self.feat_losses += [
+            self.base_loss(f_in, f_out) * w
+            for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)
+        ]
+
+        self.metrics = dict(zip(self.metric_names, self.feat_losses))
+        return sum(self.feat_losses)
+
+    def __del__(self):
+        self.hooks.remove()
+
+
+# Refactored code, originally from https://github.com/VinceMarron/style_transfer
+class WassFeatureLoss(nn.Module):
+    def __init__(self, layer_wgts=[5, 15, 2], wass_wgts=[3.0, 0.7, 0.01]):
+        super().__init__()
+        self.m_feat = models.vgg16_bn(True).features.cuda().eval()
+        requires_grad(self.m_feat, False)
+        blocks = [
+            i - 1
+            for i, o in enumerate(children(self.m_feat))
+            if isinstance(o, nn.MaxPool2d)
+        ]
+        layer_ids = blocks[2:5]
+        self.loss_features = [self.m_feat[i] for i in layer_ids]
+        self.hooks = hook_outputs(self.loss_features, detach=False)
+        self.wgts = layer_wgts
+        self.wass_wgts = wass_wgts
+        self.metric_names = (
+            ['pixel']
+            + [f'feat_{i}' for i in range(len(layer_ids))]
+            + [f'wass_{i}' for i in range(len(layer_ids))]
+        )
+        self.base_loss = F.l1_loss
+
+    def _make_features(self, x, clone=False):
+        self.m_feat(x)
+        return [(o.clone() if clone else o) for o in self.hooks.stored]
+
+    def _calc_2_moments(self, tensor):
+        chans = tensor.shape[1]
+        tensor = tensor.view(1, chans, -1)
+        n = tensor.shape[2]
+        mu = tensor.mean(2)
+        tensor = (tensor - mu[:, :, None]).squeeze(0)
+        # Prevents nasty bug that happens very occassionally- divide by zero.  Why such things happen?
+        if n == 0:
+            return None, None
+        cov = torch.mm(tensor, tensor.t()) / float(n)
+        return mu, cov
+
+    def _get_style_vals(self, tensor):
+        mean, cov = self._calc_2_moments(tensor)
+        if mean is None:
+            return None, None, None
+        eigvals, eigvects = torch.symeig(cov, eigenvectors=True)
+        eigroot_mat = torch.diag(torch.sqrt(eigvals.clamp(min=0)))
+        root_cov = torch.mm(torch.mm(eigvects, eigroot_mat), eigvects.t())
+        tr_cov = eigvals.clamp(min=0).sum()
+        return mean, tr_cov, root_cov
+
+    def _calc_l2wass_dist(
+        self, mean_stl, tr_cov_stl, root_cov_stl, mean_synth, cov_synth
+    ):
+        tr_cov_synth = torch.symeig(cov_synth, eigenvectors=True)[0].clamp(min=0).sum()
+        mean_diff_squared = (mean_stl - mean_synth).pow(2).sum()
+        cov_prod = torch.mm(torch.mm(root_cov_stl, cov_synth), root_cov_stl)
+        var_overlap = torch.sqrt(
+            torch.symeig(cov_prod, eigenvectors=True)[0].clamp(min=0) + 1e-8
+        ).sum()
+        dist = mean_diff_squared + tr_cov_stl + tr_cov_synth - 2 * var_overlap
+        return dist
+
+    def _single_wass_loss(self, pred, targ):
+        mean_test, tr_cov_test, root_cov_test = targ
+        mean_synth, cov_synth = self._calc_2_moments(pred)
+        loss = self._calc_l2wass_dist(
+            mean_test, tr_cov_test, root_cov_test, mean_synth, cov_synth
+        )
+        return loss
+
+    def forward(self, input, target):
+        out_feat = self._make_features(target, clone=True)
+        in_feat = self._make_features(input)
+        self.feat_losses = [self.base_loss(input, target)]
+        self.feat_losses += [
+            self.base_loss(f_in, f_out) * w
+            for f_in, f_out, w in zip(in_feat, out_feat, self.wgts)
+        ]
+
+        styles = [self._get_style_vals(i) for i in out_feat]
+
+        if styles[0][0] is not None:
+            self.feat_losses += [
+                self._single_wass_loss(f_pred, f_targ) * w
+                for f_pred, f_targ, w in zip(in_feat, styles, self.wass_wgts)
+            ]
+
+        self.metrics = dict(zip(self.metric_names, self.feat_losses))
+        return sum(self.feat_losses)
+
+    def __del__(self):
+        self.hooks.remove()

src/deoldify/save.py

@@ -0,0 +1,29 @@
+from fastai.basic_train import Learner, LearnerCallback
+from fastai.vision.gan import GANLearner
+
+
+class GANSaveCallback(LearnerCallback):
+    """A `LearnerCallback` that saves history of metrics while training `learn` into CSV `filename`."""
+
+    def __init__(
+        self,
+        learn: GANLearner,
+        learn_gen: Learner,
+        filename: str,
+        save_iters: int = 1000,
+    ):
+        super().__init__(learn)
+        self.learn_gen = learn_gen
+        self.filename = filename
+        self.save_iters = save_iters
+
+    def on_batch_end(self, iteration: int, epoch: int, **kwargs) -> None:
+        if iteration == 0:
+            return
+
+        if iteration % self.save_iters == 0:
+            self._save_gen_learner(iteration=iteration, epoch=epoch)
+
+    def _save_gen_learner(self, iteration: int, epoch: int):
+        filename = '{}_{}_{}'.format(self.filename, epoch, iteration)
+        self.learn_gen.save(filename)

src/deoldify/unet.py

@@ -0,0 +1,285 @@
+from fastai.layers import *
+from .layers import *
+from fastai.torch_core import *
+from fastai.callbacks.hooks import *
+from fastai.vision import *
+
+
+# The code below is meant to be merged into fastaiv1 ideally
+
+__all__ = ['DynamicUnetDeep', 'DynamicUnetWide']
+
+
+def _get_sfs_idxs(sizes: Sizes) -> List[int]:
+    "Get the indexes of the layers where the size of the activation changes."
+    feature_szs = [size[-1] for size in sizes]
+    sfs_idxs = list(
+        np.where(np.array(feature_szs[:-1]) != np.array(feature_szs[1:]))[0]
+    )
+    if feature_szs[0] != feature_szs[1]:
+        sfs_idxs = [0] + sfs_idxs
+    return sfs_idxs
+
+
+class CustomPixelShuffle_ICNR(nn.Module):
+    "Upsample by `scale` from `ni` filters to `nf` (default `ni`), using `nn.PixelShuffle`, `icnr` init, and `weight_norm`."
+
+    def __init__(
+        self,
+        ni: int,
+        nf: int = None,
+        scale: int = 2,
+        blur: bool = False,
+        leaky: float = None,
+        **kwargs
+    ):
+        super().__init__()
+        nf = ifnone(nf, ni)
+        self.conv = custom_conv_layer(
+            ni, nf * (scale ** 2), ks=1, use_activ=False, **kwargs
+        )
+        icnr(self.conv[0].weight)
+        self.shuf = nn.PixelShuffle(scale)
+        # Blurring over (h*w) kernel
+        # "Super-Resolution using Convolutional Neural Networks without Any Checkerboard Artifacts"
+        # - https://arxiv.org/abs/1806.02658
+        self.pad = nn.ReplicationPad2d((1, 0, 1, 0))
+        self.blur = nn.AvgPool2d(2, stride=1)
+        self.relu = relu(True, leaky=leaky)
+
+    def forward(self, x):
+        x = self.shuf(self.relu(self.conv(x)))
+        return self.blur(self.pad(x)) if self.blur else x
+
+
+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
+    ):
+        super().__init__()
+        self.hook = hook
+        self.shuf = CustomPixelShuffle_ICNR(
+            up_in_c, up_in_c // 2, blur=blur, leaky=leaky, **kwargs
+        )
+        self.bn = batchnorm_2d(x_in_c)
+        ni = up_in_c // 2 + x_in_c
+        nf = int((ni if final_div else ni // 2) * nf_factor)
+        self.conv1 = custom_conv_layer(ni, nf, leaky=leaky, **kwargs)
+        self.conv2 = custom_conv_layer(
+            nf, nf, leaky=leaky, self_attention=self_attention, **kwargs
+        )
+        self.relu = relu(leaky=leaky)
+
+    def forward(self, up_in: Tensor) -> Tensor:
+        s = self.hook.stored
+        up_out = self.shuf(up_in)
+        ssh = s.shape[-2:]
+        if ssh != up_out.shape[-2:]:
+            up_out = F.interpolate(up_out, s.shape[-2:], mode='nearest')
+        cat_x = self.relu(torch.cat([up_out, self.bn(s)], dim=1))
+        return self.conv2(self.conv1(cat_x))
+
+
+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,
+        norm_type: Optional[NormType] = NormType.Batch,
+        nf_factor: float = 1.0,
+        **kwargs
+    ):
+        extra_bn = norm_type == NormType.Spectral
+        imsize = (256, 256)
+        sfs_szs = model_sizes(encoder, size=imsize)
+        sfs_idxs = list(reversed(_get_sfs_idxs(sfs_szs)))
+        self.sfs = hook_outputs([encoder[i] for i in sfs_idxs], detach=False)
+        x = dummy_eval(encoder, imsize).detach()
+
+        ni = sfs_szs[-1][1]
+        middle_conv = nn.Sequential(
+            custom_conv_layer(
+                ni, ni * 2, norm_type=norm_type, extra_bn=extra_bn, **kwargs
+            ),
+            custom_conv_layer(
+                ni * 2, ni, norm_type=norm_type, extra_bn=extra_bn, **kwargs
+            ),
+        ).eval()
+        x = middle_conv(x)
+        layers = [encoder, batchnorm_2d(ni), nn.ReLU(), middle_conv]
+
+        for i, idx in enumerate(sfs_idxs):
+            not_final = i != len(sfs_idxs) - 1
+            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 = 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)
+
+        ni = x.shape[1]
+        if imsize != sfs_szs[0][-2:]:
+            layers.append(PixelShuffle_ICNR(ni, **kwargs))
+        if last_cross:
+            layers.append(MergeLayer(dense=True))
+            ni += in_channels(encoder)
+            layers.append(res_block(ni, bottle=bottle, norm_type=norm_type, **kwargs))
+        layers += [
+            custom_conv_layer(ni, n_classes, ks=1, use_activ=False, norm_type=norm_type)
+        ]
+        if y_range is not None:
+            layers.append(SigmoidRange(*y_range))
+        super().__init__(*layers)
+
+    def __del__(self):
+        if hasattr(self, "sfs"):
+            self.sfs.remove()
+
+
+# ------------------------------------------------------
+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
+    ):
+        super().__init__()
+        self.hook = hook
+        up_out = x_out = n_out // 2
+        self.shuf = CustomPixelShuffle_ICNR(
+            up_in_c, up_out, blur=blur, leaky=leaky, **kwargs
+        )
+        self.bn = batchnorm_2d(x_in_c)
+        ni = up_out + x_in_c
+        self.conv = custom_conv_layer(
+            ni, x_out, leaky=leaky, self_attention=self_attention, **kwargs
+        )
+        self.relu = relu(leaky=leaky)
+
+    def forward(self, up_in: Tensor) -> Tensor:
+        s = self.hook.stored
+        up_out = self.shuf(up_in)
+        ssh = s.shape[-2:]
+        if ssh != up_out.shape[-2:]:
+            up_out = F.interpolate(up_out, s.shape[-2:], mode='nearest')
+        cat_x = self.relu(torch.cat([up_out, self.bn(s)], dim=1))
+        return self.conv(cat_x)
+
+
+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,
+        norm_type: Optional[NormType] = NormType.Batch,
+        nf_factor: int = 1,
+        **kwargs
+    ):
+
+        nf = 512 * nf_factor
+        extra_bn = norm_type == NormType.Spectral
+        imsize = (256, 256)
+        sfs_szs = model_sizes(encoder, size=imsize)
+        sfs_idxs = list(reversed(_get_sfs_idxs(sfs_szs)))
+        self.sfs = hook_outputs([encoder[i] for i in sfs_idxs], detach=False)
+        x = dummy_eval(encoder, imsize).detach()
+
+        ni = sfs_szs[-1][1]
+        middle_conv = nn.Sequential(
+            custom_conv_layer(
+                ni, ni * 2, norm_type=norm_type, extra_bn=extra_bn, **kwargs
+            ),
+            custom_conv_layer(
+                ni * 2, ni, norm_type=norm_type, extra_bn=extra_bn, **kwargs
+            ),
+        ).eval()
+        x = middle_conv(x)
+        layers = [encoder, batchnorm_2d(ni), nn.ReLU(), middle_conv]
+
+        for i, idx in enumerate(sfs_idxs):
+            not_final = i != len(sfs_idxs) - 1
+            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)
+
+            n_out = nf if not_final else nf // 2
+
+            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)
+
+        ni = x.shape[1]
+        if imsize != sfs_szs[0][-2:]:
+            layers.append(PixelShuffle_ICNR(ni, **kwargs))
+        if last_cross:
+            layers.append(MergeLayer(dense=True))
+            ni += in_channels(encoder)
+            layers.append(res_block(ni, bottle=bottle, norm_type=norm_type, **kwargs))
+        layers += [
+            custom_conv_layer(ni, n_classes, ks=1, use_activ=False, norm_type=norm_type)
+        ]
+        if y_range is not None:
+            layers.append(SigmoidRange(*y_range))
+        super().__init__(*layers)
+
+    def __del__(self):
+        if hasattr(self, "sfs"):
+            self.sfs.remove()

src/deoldify/visualize.py

@@ -0,0 +1,247 @@
+import cv2
+import gc
+import requests
+from io import BytesIO
+import base64
+from scipy import misc
+from PIL import Image
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
+from typing import Tuple
+
+import torch
+from fastai.core import *
+from fastai.vision import *
+
+from .filters import IFilter, MasterFilter, ColorizerFilter
+from .generators import gen_inference_deep, gen_inference_wide
+
+
+
+# class LoadedModel
+class ModelImageVisualizer:
+    def __init__(self, filter: IFilter, results_dir: str = None):
+        self.filter = filter
+        self.results_dir = None if results_dir is None else Path(results_dir)
+        self.results_dir.mkdir(parents=True, exist_ok=True)
+
+    def _clean_mem(self):
+        torch.cuda.empty_cache()
+        # gc.collect()
+
+    def _open_pil_image(self, path: Path) -> Image:
+        return Image.open(path).convert('RGB')
+
+    def _get_image_from_url(self, url: str) -> Image:
+        response = requests.get(url, timeout=30, headers={'Accept': '*/*;q=0.8'})
+        img = Image.open(BytesIO(response.content)).convert('RGB')
+        return img
+
+    def plot_transformed_image_from_url(
+        self,
+        url: str,
+        path: str = 'test_images/image.png',
+        results_dir:Path = None,
+        figsize: Tuple[int, int] = (20, 20),
+        render_factor: int = None,
+        
+        display_render_factor: bool = False,
+        compare: bool = False,
+        post_process: bool = True,
+        watermarked: bool = True,
+    ) -> Path:
+        img = self._get_image_from_url(url)
+        img.save(path)
+        return self.plot_transformed_image(
+            path=path,
+            results_dir=results_dir,
+            figsize=figsize,
+            render_factor=render_factor,
+            display_render_factor=display_render_factor,
+            compare=compare,
+            post_process = post_process,
+            watermarked=watermarked,
+        )
+
+    def plot_transformed_image(
+        self,
+        path: str,
+        results_dir:Path = None,
+        figsize: Tuple[int, int] = (20, 20),
+        render_factor: int = None,
+        display_render_factor: bool = False,
+        compare: bool = False,
+        post_process: bool = True,
+        watermarked: bool = True,
+    ) -> Path:
+        path = Path(path)
+        if results_dir is None:
+            results_dir = Path(self.results_dir)
+        result = self.get_transformed_image(
+            path, render_factor, post_process=post_process,watermarked=watermarked
+        )
+        orig = self._open_pil_image(path)
+        if compare:
+            self._plot_comparison(
+                figsize, render_factor, display_render_factor, orig, result
+            )
+        else:
+            self._plot_solo(figsize, render_factor, display_render_factor, result)
+
+        orig.close()
+        result_path = self._save_result_image(path, result, results_dir=results_dir)
+        result.close()
+        return result_path
+
+    def plot_transformed_pil_image(
+        self,
+        input_image: Image,
+        figsize: Tuple[int, int] = (20, 20),
+        render_factor: int = None,
+        display_render_factor: bool = False,
+        compare: bool = False,
+        post_process: bool = True,
+    ) -> Image:
+
+        result = self.get_transformed_pil_image(
+            input_image, render_factor, post_process=post_process
+        )
+
+        if compare:
+            self._plot_comparison(
+                figsize, render_factor, display_render_factor, input_image, result
+            )
+        else:
+            self._plot_solo(figsize, render_factor, display_render_factor, result)
+
+        return result
+
+    def _plot_comparison(
+        self,
+        figsize: Tuple[int, int],
+        render_factor: int,
+        display_render_factor: bool,
+        orig: Image,
+        result: Image,
+    ):
+        fig, axes = plt.subplots(1, 2, figsize=figsize)
+        self._plot_image(
+            orig,
+            axes=axes[0],
+            figsize=figsize,
+            render_factor=render_factor,
+            display_render_factor=False,
+        )
+        self._plot_image(
+            result,
+            axes=axes[1],
+            figsize=figsize,
+            render_factor=render_factor,
+            display_render_factor=display_render_factor,
+        )
+
+    def _plot_solo(
+        self,
+        figsize: Tuple[int, int],
+        render_factor: int,
+        display_render_factor: bool,
+        result: Image,
+    ):
+        fig, axes = plt.subplots(1, 1, figsize=figsize)
+        self._plot_image(
+            result,
+            axes=axes,
+            figsize=figsize,
+            render_factor=render_factor,
+            display_render_factor=display_render_factor,
+        )
+
+    def _save_result_image(self, source_path: Path, image: Image, results_dir = None) -> Path:
+        if results_dir is None:
+            results_dir = Path(self.results_dir)
+        result_path = results_dir / source_path.name
+        image.save(result_path)
+        return result_path
+
+    def get_transformed_image(
+        self, path: Path, render_factor: int = None, post_process: bool = True,
+        watermarked: bool = True,
+    ) -> Image:
+        self._clean_mem()
+        orig_image = self._open_pil_image(path)
+        filtered_image = self.filter.filter(
+            orig_image, orig_image, render_factor=render_factor,post_process=post_process
+        )
+
+        return filtered_image
+
+    def get_transformed_pil_image(
+        self, input_image: Image, render_factor: int = None, post_process: bool = True,
+    ) -> Image:
+        self._clean_mem()
+        filtered_image = self.filter.filter(
+            input_image, input_image, render_factor=render_factor,post_process=post_process
+        )
+
+        return filtered_image
+
+    def _plot_image(
+        self,
+        image: Image,
+        render_factor: int,
+        axes: Axes = None,
+        figsize=(20, 20),
+        display_render_factor = False,
+    ):
+        if axes is None:
+            _, axes = plt.subplots(figsize=figsize)
+        axes.imshow(np.asarray(image) / 255)
+        axes.axis('off')
+        if render_factor is not None and display_render_factor:
+            plt.text(
+                10,
+                10,
+                'render_factor: ' + str(render_factor),
+                color='white',
+                backgroundcolor='black',
+            )
+
+    def _get_num_rows_columns(self, num_images: int, max_columns: int) -> Tuple[int, int]:
+        columns = min(num_images, max_columns)
+        rows = num_images // columns
+        rows = rows if rows * columns == num_images else rows + 1
+        return rows, columns
+
+
+def get_image_colorizer(
+    root_folder: Path = Path('./'), render_factor: int = 35, artistic: bool = True
+) -> ModelImageVisualizer:
+    if artistic:
+        return get_artistic_image_colorizer(root_folder=root_folder, render_factor=render_factor)
+    else:
+        return get_stable_image_colorizer(root_folder=root_folder, render_factor=render_factor)
+
+
+def get_stable_image_colorizer(
+    root_folder: Path = Path('./'),
+    weights_name: str = 'ColorizeStable_gen',
+    results_dir='output',
+    render_factor: int = 35
+) -> ModelImageVisualizer:
+    learn = gen_inference_wide(root_folder=root_folder, weights_name=weights_name)
+    filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
+    vis = ModelImageVisualizer(filtr, results_dir=results_dir)
+    return vis
+
+
+def get_artistic_image_colorizer(
+    root_folder: Path = Path('./'),
+    weights_name: str = 'ColorizeArtistic_gen',
+    results_dir='output',
+    render_factor: int = 35
+) -> ModelImageVisualizer:
+    learn = gen_inference_deep(root_folder=root_folder, weights_name=weights_name)
+    filtr = MasterFilter([ColorizerFilter(learn=learn)], render_factor=render_factor)
+    vis = ModelImageVisualizer(filtr, results_dir=results_dir)
+    return vis

src/st_style.py

@@ -0,0 +1,42 @@
+button_style = """
+<style>
+div.stButton > button:first-child {
+    background-color: rgb(255, 75, 75);
+    color: rgb(255, 255, 255);
+}
+div.stButton > button:hover {
+    background-color: rgb(255, 75, 75);
+    color: rgb(255, 255, 255);
+}
+div.stButton > button:active {
+    background-color: rgb(255, 75, 75);
+    color: rgb(255, 255, 255);
+}
+div.stButton > button:focus {
+    background-color: rgb(255, 75, 75);
+    color: rgb(255, 255, 255);
+}
+.css-1cpxqw2:focus:not(:active) {
+    background-color: rgb(255, 75, 75);
+    border-color: rgb(255, 75, 75);
+    color: rgb(255, 255, 255);
+}
+"""
+
+style = """
+<style>
+#MainMenu {
+    visibility: hidden;
+}
+footer {
+    visibility: hidden;
+}
+header {
+    visibility: hidden;
+}
+</style>
+"""
+
+
+def apply_prod_style(st):
+    return st.markdown(style, unsafe_allow_html=True)