repo_description: Deepfakes Software For All
repo_topics: deep-face-swap deep-learning deep-neural-networks deepface deepfakes deeplearning face-swap faceswap fakeapp machine-learning myfakeapp neural-nets neural-networks openfaceswap
code: #!/usr/bin/env python3
  """
  Library for serializing python objects to and from various different serializer formats
  """
  
  import json
  import logging
  import os
  import pickle
  import zlib
  
  from io import BytesIO
  
  import numpy as np
  
  from lib.utils import FaceswapError, get_module_objects
  
  try:
      import yaml
      _HAS_YAML = True
  except ImportError:
      _HAS_YAML = False
  
  logger = logging.getLogger(__name__)
  
  
  class Serializer():
      """ A convenience class for various serializers.
  
      This class should not be called directly as it acts as the parent for various serializers.
      All serializers should be called from :func:`get_serializer` or
      :func:`get_serializer_from_filename`
  
      Example
      -------
      >>> from lib.serializer import get_serializer
      >>> serializer = get_serializer('json')
      >>> json_file = '/path/to/json/file.json'
      >>> data = serializer.load(json_file)
      >>> serializer.save(json_file, data)
  
      """
      def __init__(self):
          self._file_extension = None
          self._write_option = "wb"
          self._read_option = "rb"
  
      @property
      def file_extension(self):
          """ str: The file extension of the serializer """
          return self._file_extension
  
      def save(self, filename, data):
          """ Serialize data and save to a file
  
          Parameters
          ----------
          filename: str
              The path to where the serialized file should be saved
          data: varies
              The data that is to be serialized to file
  
          Example
          ------
          >>> serializer = get_serializer('json')
          >>> data ['foo', 'bar']
          >>> json_file = '/path/to/json/file.json'
          >>> serializer.save(json_file, data)
          """
          logger.debug("filename: %s, data type: %s", filename, type(data))
          filename = self._check_extension(filename)
          try:
              with open(filename, self._write_option) as s_file:
                  s_file.write(self.marshal(data))
          except IOError as err:
              msg = f"Error writing to '{filename}': {err.strerror}"
              raise FaceswapError(msg) from err
  
      def _check_extension(self, filename):
          """ Check the filename has an extension. If not add the correct one for the serializer """
          extension = os.path.splitext(filename)[1]
          retval = filename if extension else f"{filename}.{self.file_extension}"
          logger.debug("Original filename: '%s', final filename: '%s'", filename, retval)
          return retval
  
      def load(self, filename):
          """ Load data from an existing serialized file
  
          Parameters
          ----------
          filename: str
              The path to the serialized file
  
          Returns
          ----------
          data: varies
              The data in a python object format
  
          Example
          ------
          >>> serializer = get_serializer('json')
          >>> json_file = '/path/to/json/file.json'
          >>> data = serializer.load(json_file)
          """
          logger.debug("filename: %s", filename)
          try:
              with open(filename, self._read_option) as s_file:
                  data = s_file.read()
                  logger.debug("stored data type: %s", type(data))
                  retval = self.unmarshal(data)
  
          except IOError as err:
              msg = f"Error reading from '{filename}': {err.strerror}"
              raise FaceswapError(msg) from err
          logger.debug("data type: %s", type(retval))
          return retval
  
      def marshal(self, data):
          """ Serialize an object
  
          Parameters
          ----------
          data: varies
              The data that is to be serialized
  
          Returns
          -------
          data: varies
              The data in a the serialized data format
  
          Example
          ------
          >>> serializer = get_serializer('json')
          >>> data ['foo', 'bar']
          >>> json_data = serializer.marshal(data)
          """
          logger.debug("data type: %s", type(data))
          try:
              retval = self._marshal(data)
          except Exception as err:
              msg = f"Error serializing data for type {type(data)}: {str(err)}"
              raise FaceswapError(msg) from err
          logger.debug("returned data type: %s", type(retval))
          return retval
  
      def unmarshal(self, serialized_data):
          """ Unserialize data to its original object type
  
          Parameters
          ----------
          serialized_data: varies
              Data in serializer format that is to be unmarshalled to its original object
  
          Returns
          -------
          data: varies
              The data in a python object format
  
          Example
          ------
          >>> serializer = get_serializer('json')
          >>> json_data = <json object>
          >>> data = serializer.unmarshal(json_data)
          """
          logger.debug("data type: %s", type(serialized_data))
          try:
              retval = self._unmarshal(serialized_data)
          except Exception as err:
              msg = f"Error unserializing data for type {type(serialized_data)}: {str(err)}"
              raise FaceswapError(msg) from err
          logger.debug("returned data type: %s", type(retval))
          return retval
  
      def _marshal(self, data):
          """ Override for serializer specific marshalling """
          raise NotImplementedError()
  
      def _unmarshal(self, data):
          """ Override for serializer specific unmarshalling """
          raise NotImplementedError()
  
  
  class _YAMLSerializer(Serializer):
      """ YAML Serializer """
      def __init__(self):
          super().__init__()
          self._file_extension = "yml"
  
      def _marshal(self, data):
          return yaml.dump(data, default_flow_style=False).encode("utf-8")
  
      def _unmarshal(self, data):
          return yaml.load(data.decode("utf-8", errors="replace"), Loader=yaml.FullLoader)
  
  
  class _JSONSerializer(Serializer):
      """ JSON Serializer """
      def __init__(self):
          super().__init__()
          self._file_extension = "json"
  
      def _marshal(self, data):
          return json.dumps(data, indent=2).encode("utf-8")
  
      def _unmarshal(self, data):
          return json.loads(data.decode("utf-8", errors="replace"))
  
  
  class _PickleSerializer(Serializer):
      """ Pickle Serializer """
      def __init__(self):
          super().__init__()
          self._file_extension = "pickle"
  
      def _marshal(self, data):
          return pickle.dumps(data)
  
      def _unmarshal(self, data):
          return pickle.loads(data)
  
  
  class _NPYSerializer(Serializer):
      """ NPY Serializer """
      def __init__(self):
          super().__init__()
          self._file_extension = "npy"
          self._bytes = BytesIO()
  
      def _marshal(self, data):
          """ NPY Marshal to bytesIO so standard bytes writer can write out """
          b_handler = BytesIO()
          np.save(b_handler, data)
          b_handler.seek(0)
          return b_handler.read()
  
      def _unmarshal(self, data):
          """ NPY Unmarshal to bytesIO so we can use numpy loader """
          b_handler = BytesIO(data)
          retval = np.load(b_handler)
          del b_handler
          if retval.dtype == "object":
              retval = retval[()]
          return retval
  
  
  class _CompressedSerializer(Serializer):
      """ A compressed pickle serializer for Faceswap """
      def __init__(self):
          super().__init__()
          self._file_extension = "fsa"
          self._child = get_serializer("pickle")
  
      def _marshal(self, data):
          """ Pickle and compress data """
          data = self._child._marshal(data)  # pylint:disable=protected-access
          return zlib.compress(data)
  
      def _unmarshal(self, data):
          """ Decompress and unpicke data """
          data = zlib.decompress(data)
          return self._child._unmarshal(data)  # pylint:disable=protected-access
  
  
  def get_serializer(serializer):
      """ Obtain a serializer object
  
      Parameters
      ----------
      serializer: {'json', 'pickle', yaml', 'npy', 'compressed'}
          The required serializer format
  
      Returns
      -------
      serializer: :class:`Serializer`
          A serializer object for handling the requested data format
  
      Example
      -------
      >>> serializer = get_serializer('json')
      """
      retval = None
      if serializer.lower() == "npy":
          retval = _NPYSerializer()
      elif serializer.lower() == "compressed":
          retval = _CompressedSerializer()
      elif serializer.lower() == "json":
          retval = _JSONSerializer()
      elif serializer.lower() == "pickle":
          retval = _PickleSerializer()
      elif serializer.lower() == "yaml" and _HAS_YAML:
          retval = _YAMLSerializer()
      elif serializer.lower() == "yaml":
          logger.warning("You must have PyYAML installed to use YAML as the serializer."
                         "Switching to JSON as the serializer.")
          retval = _JSONSerializer
      else:
          logger.warning("Unrecognized serializer: '%s'. Returning json serializer", serializer)
      logger.debug(retval)
      return retval
  
  
  def get_serializer_from_filename(filename):
      """ Obtain a serializer object from a filename
  
      Parameters
      ----------
      filename: str
          Filename to determine the serializer type from
  
      Returns
      -------
      serializer: :class:`Serializer`
          A serializer object for handling the requested data format
  
      Example
      -------
      >>> filename = '/path/to/json/file.json'
      >>> serializer = get_serializer_from_filename(filename)
      """
      logger.debug("filename: '%s'", filename)
      extension = os.path.splitext(filename)[1].lower()
      logger.debug("extension: '%s'", extension)
  
      if extension == ".json":
          retval = _JSONSerializer()
      elif extension in (".p", ".pickle"):
          retval = _PickleSerializer()
      elif extension == ".npy":
          retval = _NPYSerializer()
      elif extension == ".fsa":
          retval = _CompressedSerializer()
      elif extension in (".yaml", ".yml") and _HAS_YAML:
          retval = _YAMLSerializer()
      elif extension in (".yaml", ".yml"):
          logger.warning("You must have PyYAML installed to use YAML as the serializer.\n"
                         "Switching to JSON as the serializer.")
          retval = _JSONSerializer()
      else:
          logger.warning("Unrecognized extension: '%s'. Returning json serializer", extension)
          retval = _JSONSerializer()
      logger.debug(retval)
      return retval
  
  
  __all__ = get_module_objects(__name__)
  