DeOldify/fasterai/filters.py

from numpy import ndarray
from abc import ABC, abstractmethod
from .generators import Unet34, Unet101, GeneratorModule
from .transforms import BlackAndWhiteTransform
from fastai.torch_imports import *
from fastai.core import *
from fastai.transforms import Transform, scale_min, tfms_from_stats, inception_stats
from fastai.transforms import CropType, NoCrop, Denormalize, Scale, scale_to
import math
from scipy import misc

class Padding():
    def __init__(self, top:int, bottom:int, left:int, right:int):
        self.top = top
        self.bottom = bottom
        self.left = left
        self.right = right
  
class Filter(ABC):
    def __init__(self, tfms:[Transform]):
        super().__init__()
        self.tfms=tfms
        self.denorm = Denormalize(*inception_stats)
    
    @abstractmethod
    def filter(self, orig_image:ndarray, filtered_image:ndarray, render_factor:int)->ndarray:
        pass

    def _init_model(self, model:nn.Module, weights_path:Path):
        load_model(model, weights_path)
        model.eval()
        torch.no_grad() 

    def _transform(self, orig:ndarray, sz:int):
        for tfm in self.tfms:
            orig,_=tfm(orig, False)
        _,val_tfms = tfms_from_stats(inception_stats, sz, crop_type=CropType.NO, aug_tfms=[])
        val_tfms.tfms = [tfm for tfm in val_tfms.tfms if not (isinstance(tfm, NoCrop) or isinstance(tfm, Scale))]
        orig = val_tfms(orig)
        return orig

    def _scale_to_square(self, orig:ndarray, targ:int, interpolation=cv2.INTER_AREA):
        r,c,*_ = orig.shape
        ratio = targ/max(r,c)
        #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!
        sz = (targ, targ)
        return cv2.resize(orig, sz, interpolation=interpolation)

    def _get_model_ready_image_ndarray(self, orig:ndarray, sz:int):
        result = self._scale_to_square(orig, sz)
        sz=result.shape[0]
        result = self._transform(result, sz)
        return result

    def _denorm(self, image: ndarray):
        if len(image.shape)==3: 
            image = image[None]
        return self.denorm(np.rollaxis(image,1,4))

    def _model_process(self, model:GeneratorModule, orig:ndarray, sz:int, gpu:int):
        orig = self._get_model_ready_image_ndarray(orig, sz)
        orig = VV_(orig[None]) 
        orig = orig.to(device=gpu)
        result = model(orig)
        result = result.detach().cpu().numpy()
        result = self._denorm(result)
        return result[0]

    def _convert_to_pil(self, im_array:ndarray):
        im_array = np.clip(im_array,0,1)
        return misc.toimage(im_array)

    def _unsquare(self, result:ndarray, orig:ndarray):
        sz = (orig.shape[1], orig.shape[0])
        return cv2.resize(result, sz, interpolation=cv2.INTER_AREA)  



class AbstractColorizer(Filter):
    def __init__(self, gpu:int, weights_path:Path, nf_factor:int=2, map_to_orig:bool=True):
        super().__init__(tfms=[BlackAndWhiteTransform()])
        self.model = self._get_model(nf_factor=nf_factor, gpu=gpu)
        self.gpu = gpu
        self._init_model(self.model, weights_path)
        self.render_base=16
        self.map_to_orig=map_to_orig

    @abstractmethod
    def _get_model(self, nf_factor:int, gpu:int)->GeneratorModule:
        pass
    
    def filter(self, orig_image:ndarray, filtered_image:ndarray, render_factor:int=36)->ndarray:
        render_sz = render_factor * self.render_base
        model_image = self._model_process(self.model, orig=filtered_image, sz=render_sz, gpu=self.gpu)
        if self.map_to_orig:
            return self._post_process(model_image, orig_image)
        else:
            return self._post_process(model_image, filtered_image)


    #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:ndarray, orig:ndarray):
        for tfm in self.tfms:
            orig,_=tfm(orig, False)

        raw_color = self._unsquare(raw_color, orig)
        color_yuv = cv2.cvtColor(raw_color, cv2.COLOR_BGR2YUV)
        #do a black and white transform first to get better luminance values
        orig_yuv = cv2.cvtColor(orig, cv2.COLOR_BGR2YUV)
        hires = np.copy(orig_yuv)
        hires[:,:,1:3] = color_yuv[:,:,1:3]
        return cv2.cvtColor(hires, cv2.COLOR_YUV2BGR)   

class Colorizer34(AbstractColorizer):
    def __init__(self, gpu:int, weights_path:Path, nf_factor:int=2, map_to_orig:bool=True):
        super().__init__(gpu=gpu, weights_path=weights_path, nf_factor=nf_factor, map_to_orig=map_to_orig)

    def _get_model(self, nf_factor:int, gpu:int)->GeneratorModule:
        return Unet34(nf_factor=nf_factor).cuda(gpu)


class Colorizer101(AbstractColorizer):
    def __init__(self, gpu:int, weights_path:Path, nf_factor:int=2, map_to_orig:bool=True):
        super().__init__(gpu=gpu, weights_path=weights_path, nf_factor=nf_factor, map_to_orig=map_to_orig)

    def _get_model(self, nf_factor:int, gpu:int)->GeneratorModule:
        return Unet101(nf_factor=nf_factor).cuda(gpu)


#TODO:  May not want to do square rendering here like in colorization- it definitely loses 
#fidelity visibly (but not too terribly).  Will revisit.
class DeFader(Filter): 
    def __init__(self, gpu:int, weights_path:Path, nf_factor:int=2):
        super().__init__(tfms=[BlackAndWhiteTransform()])
        self.model = Unet34(nf_factor=nf_factor).cuda(gpu)
        self._init_model(self.model, weights_path)
        self.render_base=16
        self.gpu = gpu

    def filter(self, orig_image:ndarray, filtered_image:ndarray, render_factor:int=36)->ndarray:
        render_sz = render_factor * self.render_base
        model_image = self._model_process(self.model, orig=filtered_image, sz=render_sz, gpu=self.gpu)
        return self._post_process(model_image, filtered_image)

    def _post_process(self, result:ndarray, orig:ndarray):
        return self._unsquare(result, orig)