first commit in hf_spaces

app.py

@@ -0,0 +1,160 @@
+#importing the libraries
+import os, sys, re
+import streamlit as st
+from PIL import Image
+import cv2
+import numpy as np
+import uuid
+
+# Import torch libraries
+import fastai
+import torch
+
+# Import util functions from app_utils
+from app_utils import download
+from app_utils import generate_random_filename
+from app_utils import clean_me
+from app_utils import clean_all
+from app_utils import create_directory
+from app_utils import get_model_bin
+from app_utils import convertToJPG
+
+# Import util functions from deoldify
+# NOTE:  This must be the first call in order to work properly!
+from deoldify import device
+from deoldify.device_id import DeviceId
+#choices:  CPU, GPU0...GPU7
+device.set(device=DeviceId.CPU)
+from deoldify.visualize import *
+
+
+####### INPUT PARAMS ###########
+model_folder = 'models/'
+max_img_size = 800
+################################
+
+@st.cache(allow_output_mutation=True)
+def load_model(model_dir, option):
+    if option.lower() == 'artistic':
+        model_url = 'https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth'
+        get_model_bin(model_url, os.path.join(model_dir, "ColorizeArtistic_gen.pth"))
+        colorizer = get_image_colorizer(artistic=True)
+    elif option.lower() == 'stable':
+        model_url = "https://www.dropbox.com/s/usf7uifrctqw9rl/ColorizeStable_gen.pth?dl=0"
+        get_model_bin(model_url, os.path.join(model_dir, "ColorizeStable_gen.pth"))
+        colorizer = get_image_colorizer(artistic=False)
+
+    return colorizer
+
+def resize_img(input_img, max_size):
+    img = input_img.copy()
+    img_height, img_width = img.shape[0],img.shape[1]
+
+    if max(img_height, img_width) > max_size:
+        if img_height > img_width:
+            new_width = img_width*(max_size/img_height)
+            new_height = max_size
+            resized_img = cv2.resize(img,(int(new_width), int(new_height)))
+            return resized_img
+
+        elif img_height <= img_width:
+            new_width = img_height*(max_size/img_width)
+            new_height = max_size
+            resized_img = cv2.resize(img,(int(new_width), int(new_height)))
+            return resized_img
+
+    return img
+
+def get_image_download_link(img,filename,text):
+    button_uuid = str(uuid.uuid4()).replace('-', '')
+    button_id = re.sub('\d+', '', button_uuid)
+    
+    custom_css = f""" 
+        <style>
+            #{button_id} {{
+                background-color: rgb(255, 255, 255);
+                color: rgb(38, 39, 48);
+                padding: 0.25em 0.38em;
+                position: relative;
+                text-decoration: none;
+                border-radius: 4px;
+                border-width: 1px;
+                border-style: solid;
+                border-color: rgb(230, 234, 241);
+                border-image: initial;
+
+            }} 
+            #{button_id}:hover {{
+                border-color: rgb(246, 51, 102);
+                color: rgb(246, 51, 102);
+            }}
+            #{button_id}:active {{
+                box-shadow: none;
+                background-color: rgb(246, 51, 102);
+                color: white;
+                }}
+        </style> """
+
+    buffered = BytesIO()
+    img.save(buffered, format="JPEG")
+    img_str = base64.b64encode(buffered.getvalue()).decode()
+    href =  custom_css + f'<a href="data:file/txt;base64,{img_str}" id="{button_id}" download="{filename}">{text}</a>'
+    return href
+
+
+# General configuration
+st.set_page_config(layout="centered")
+st.set_option('deprecation.showfileUploaderEncoding', False)
+st.markdown('''
+<style>
+    .uploadedFile {display: none}
+<style>''',
+unsafe_allow_html=True)
+
+# Main window configuration
+st.title("Black and white colorizer")
+st.markdown("This app puts color into your black and white pictures")
+title_message = st.empty()
+
+title_message.markdown("**Model loading, please wait** ⌛")
+
+# # Sidebar
+color_option = st.sidebar.selectbox('Select colorizer mode',
+                                    ('Artistic', 'Stable'))
+                                    
+# st.sidebar.title('Model parameters')
+# det_conf_thres = st.sidebar.slider("Detector confidence threshold", 0.1, 0.9, value=0.5, step=0.1)
+# det_nms_thres = st.sidebar.slider("Non-maximum supression IoU", 0.1, 0.9, value=0.4, step=0.1)
+
+# Load models
+colorizer = load_model(model_folder, color_option)
+
+title_message.markdown("**To begin, please upload an image** 👇")
+
+#Choose your own image
+uploaded_file = st.file_uploader("Upload a black and white photo", type=['png', 'jpg', 'jpeg'])
+
+# show = st.image(use_column_width='auto')
+input_img_pos = st.empty()
+output_img_pos = st.empty()
+
+if uploaded_file is not None:
+    img_name = uploaded_file.name
+
+    pil_img = Image.open(uploaded_file)
+    img_rgb = np.array(pil_img)
+
+    resized_img_rgb = resize_img(img_rgb, max_img_size)
+    resized_pil_img = Image.fromarray(resized_img_rgb)
+
+    title_message.markdown("**Processing your image, please wait** ⌛")
+
+    output_pil_img = colorizer.plot_transformed_pil_image(resized_pil_img, render_factor=35, compare=False)
+    
+    title_message.markdown("**To begin, please upload an image** 👇")
+
+    # Plot images
+    input_img_pos.image(resized_pil_img, 'Input image', use_column_width=True)
+    output_img_pos.image(output_pil_img, 'Output image', use_column_width=True)
+
+    st.markdown(get_image_download_link(output_pil_img, img_name, 'Download '+img_name), unsafe_allow_html=True)

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)
+

deoldify/__init__.py

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

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

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)

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,
+    )

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('./dummy/')
+    return get_colorize_data(
+        sz=1, bs=1, crappy_path=path, good_path=path, keep_pct=0.001
+    )

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

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

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
+
+
+# -----------------------------

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)

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()

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)

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()

deoldify/visualize.py

@@ -0,0 +1,456 @@
+from fastai.core import *
+from fastai.vision import *
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
+from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
+from .filters import IFilter, MasterFilter, ColorizerFilter
+from .generators import gen_inference_deep, gen_inference_wide
+# from tensorboardX import SummaryWriter
+from scipy import misc
+from PIL import Image
+# import ffmpeg
+# import youtube_dl
+import gc
+import requests
+from io import BytesIO
+import base64
+# from IPython import display as ipythondisplay
+# from IPython.display import HTML
+# from IPython.display import Image as ipythonimage
+import cv2
+
+
+# # adapted from https://www.pyimagesearch.com/2016/04/25/watermarking-images-with-opencv-and-python/
+# def get_watermarked(pil_image: Image) -> Image:
+#     try:
+#         image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
+#         (h, w) = image.shape[:2]
+#         image = np.dstack([image, np.ones((h, w), dtype="uint8") * 255])
+#         pct = 0.05
+#         full_watermark = cv2.imread(
+#             './resource_images/watermark.png', cv2.IMREAD_UNCHANGED
+#         )
+#         (fwH, fwW) = full_watermark.shape[:2]
+#         wH = int(pct * h)
+#         wW = int((pct * h / fwH) * fwW)
+#         watermark = cv2.resize(full_watermark, (wH, wW), interpolation=cv2.INTER_AREA)
+#         overlay = np.zeros((h, w, 4), dtype="uint8")
+#         (wH, wW) = watermark.shape[:2]
+#         overlay[h - wH - 10 : h - 10, 10 : 10 + wW] = watermark
+#         # blend the two images together using transparent overlays
+#         output = image.copy()
+#         cv2.addWeighted(overlay, 0.5, output, 1.0, 0, output)
+#         rgb_image = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+#         final_image = Image.fromarray(rgb_image)
+#         return final_image
+#     except:
+#         # Don't want this to crash everything, so let's just not watermark the image for now.
+#         return pil_image
+
+
+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 PIL.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 = PIL.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: (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: (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: (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: (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: (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
+        )
+
+        # if watermarked:
+        #     return get_watermarked(filtered_image)
+
+        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) -> (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
+
+
+# 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) -> str:
+#         probe = ffmpeg.probe(str(source_path))
+#         stream_data = next(
+#             (stream for stream in probe['streams'] if stream['codec_type'] == 'video'),
+#             None,
+#         )
+#         return stream_data['avg_frame_rate']
+
+#     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),
+#             'retries': 30,
+#             'fragment-retries': 30
+#         }
+#         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, render_factor: int = None, post_process: bool = True,
+#         watermarked: bool = True,
+#     ):
+#         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), render_factor=render_factor, post_process=post_process,watermarked=watermarked
+#                 )
+#                 color_image.save(str(colorframes_folder / img))
+
+#     def _build_video(self, source_path: Path) -> Path:
+#         colorized_path = self.result_folder / (
+#             source_path.name.replace('.mp4', '_no_audio.mp4')
+#         )
+#         colorframes_folder = self.colorframes_root / (source_path.stem)
+#         colorframes_path_template = str(colorframes_folder / '%5d.jpg')
+#         colorized_path.parent.mkdir(parents=True, exist_ok=True)
+#         if colorized_path.exists():
+#             colorized_path.unlink()
+#         fps = self._get_fps(source_path)
+
+#         ffmpeg.input(
+#             str(colorframes_path_template),
+#             format='image2',
+#             vcodec='mjpeg',
+#             framerate=fps,
+#         ).output(str(colorized_path), crf=17, vcodec='libx264').run(capture_stdout=True)
+
+#         result_path = self.result_folder / source_path.name
+#         if result_path.exists():
+#             result_path.unlink()
+#         # making copy of non-audio version in case adding back audio doesn't apply or fails.
+#         shutil.copyfile(str(colorized_path), str(result_path))
+
+#         # adding back sound here
+#         audio_file = Path(str(source_path).replace('.mp4', '.aac'))
+#         if audio_file.exists():
+#             audio_file.unlink()
+
+#         os.system(
+#             'ffmpeg -y -i "'
+#             + str(source_path)
+#             + '" -vn -acodec copy "'
+#             + str(audio_file)
+#             + '"'
+#         )
+
+#         if audio_file.exists:
+#             os.system(
+#                 'ffmpeg -y -i "'
+#                 + str(colorized_path)
+#                 + '" -i "'
+#                 + str(audio_file)
+#                 + '" -shortest -c:v copy -c:a aac -b:a 256k "'
+#                 + str(result_path)
+#                 + '"'
+#             )
+#         print('Video created here: ' + str(result_path))
+#         return result_path
+
+#     def colorize_from_url(
+#         self,
+#         source_url,
+#         file_name: str,
+#         render_factor: int = None,
+#         post_process: bool = True,
+#         watermarked: bool = True,
+
+#     ) -> Path:
+#         source_path = self.source_folder / file_name
+#         self._download_video_from_url(source_url, source_path)
+#         return self._colorize_from_path(
+#             source_path, render_factor=render_factor, post_process=post_process,watermarked=watermarked
+#         )
+
+#     def colorize_from_file_name(
+#         self, file_name: str, render_factor: int = None,  watermarked: bool = True, post_process: bool = True,
+#     ) -> Path:
+#         source_path = self.source_folder / file_name
+#         return self._colorize_from_path(
+#             source_path, render_factor=render_factor,  post_process=post_process,watermarked=watermarked
+#         )
+
+#     def _colorize_from_path(
+#         self, source_path: Path, render_factor: int = None,  watermarked: bool = True, post_process: bool = True
+#     ) -> Path:
+#         if not source_path.exists():
+#             raise Exception(
+#                 'Video at path specfied, ' + str(source_path) + ' could not be found.'
+#             )
+#         self._extract_raw_frames(source_path)
+#         self._colorize_raw_frames(
+#             source_path, render_factor=render_factor,post_process=post_process,watermarked=watermarked
+#         )
+#         return self._build_video(source_path)
+
+
+# def get_video_colorizer(render_factor: int = 21) -> VideoColorizer:
+#     return get_stable_video_colorizer(render_factor=render_factor)
+
+
+# def get_artistic_video_colorizer(
+#     root_folder: Path = Path('./'),
+#     weights_name: str = 'ColorizeArtistic_gen',
+#     results_dir='result_images',
+#     render_factor: int = 35
+# ) -> VideoColorizer:
+#     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 VideoColorizer(vis)
+
+
+# def get_stable_video_colorizer(
+#     root_folder: Path = Path('./'),
+#     weights_name: str = 'ColorizeVideo_gen',
+#     results_dir='result_images',
+#     render_factor: int = 21
+# ) -> VideoColorizer:
+#     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 VideoColorizer(vis)
+
+
+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='result_images',
+    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='result_images',
+    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

requirements.txt

@@ -0,0 +1,11 @@
+scipy==1.7.1
+scikit_image==0.18.3
+streamlit==0.88.0
+requests==2.26.0
+torch==1.9.0
+torchvision==0.10.0
+matplotlib==3.4.3
+numpy==1.21.2
+opencv_python==4.5.3.56
+fastai==1.0.51
+Pillow==8.3.2