IRestorationProcess.py

import tensorflow as tf

class IRestorationProcess(tf.keras.Model):
  @staticmethod
  def getOutputSize(outputs):
    sizes = {
      'rgb': 3,
      'grayscale': 1,
      'sobel': 6,
    }
    channels = 0
    for name in outputs:
      assert name in sizes, 'Invalid output: %s' % name
      channels += sizes[name]
      continue
    return channels
  
  def __init__(self, predictions, name=None, **kwargs):
    if name is None: name = self.__class__.__name__
    super().__init__(name=name, **kwargs)
    predictions = list(predictions)
    if 'rgb' not in predictions: predictions.insert(0, 'rgb')
    self._outputs = predictions
    self._channels = IRestorationProcess.getOutputSize(predictions)

    print('[IRestorationProcess] Restorator:', self.__class__.__name__)
    return
  
  def call(self, *args, **kwargs):
    raise RuntimeError('IRestorationProcess object cannot be called directly')
  
  def forward(self, x0, xT=None, model=None):
    raise NotImplementedError()
  
  def reverse(self, value, denoiser, modelT=None, **kwargs):
    '''
    This function implements the reverse restorator process.

    value: (B, ...) is the initial value. Typically, this is a just a noise sample.
           Can be a tuple with batch shape (B, ...)
    denoiser: callable (x, t) -> x. Denoiser function or model, which takes a value and a continuous time t and returns a denoised value.
    modelT: callable (t) -> t. Model for the time parameter. If None, then the time parameter is assumed to be a continuous value in [0, 1].

    returns: (B, ...) denoised value
    '''
    raise NotImplementedError()
  
  def withExtraOutputs(self, value, **kwargs):
    # add extras to the target
    targets = [value]
    for extra in self._outputs:
      if extra in kwargs:
        targets.append(kwargs[extra])
      continue
    target = tf.concat(targets, axis=-1)
    return target
  
  def calculate_loss(self, x_hat, predicted, **kwargs):
    raise NotImplementedError()
  
  def train_step(self, x0, model, xT=None, **kwargs):
    x_hat = self.forward(
      x0=x0, xT=xT,
      model=lambda **kwargs: tf.stop_gradient(model(**kwargs)) # stop gradient for the model
    )
    values = model(T=x_hat['T'], V=x_hat['xT'], extras=x_hat.get('extras', None))

    # we want calculate the loss WITH the residual
    totalLoss = self.calculate_loss(x_hat, values, **kwargs)
    return dict(
      loss=totalLoss,
      value=self.targets(x_hat, values)
    )
    
  def targets(self, x_hat, values):
    '''
    This function calculates the target values.
    x_hat: dictionary with the keys 'T' and 'xT'
    values: predicted values
    '''
    return values[:, :self.predictions]
  
  @property
  def channels(self):
    return self._channels
  
  @property
  def predictions(self):
    return self.getOutputSize(['rgb'])