code: def preserve_shape(func):
      
      @wraps(func)
      def wrapped_function(img, *args, **kwargs):
          shape = img.shape
          result = func(img, *args, **kwargs)
          result = result.reshape(shape)
          return result
  
      return wrapped_function
docs: Preserve shape of the image.