repo_description: DeepFaceLab is the leading software for creating deepfakes.
repo_topics: arxiv creating-deepfakes deep-face-swap deep-learning deep-neural-networks deepface deepfacelab deepfakes deeplearning face-swap faceswap fakeapp machine-learning neural-nets neural-networks
code: import numpy as np
  import cv2
  from enum import IntEnum
  
  
  class SegIEPolyType(IntEnum):
      EXCLUDE = 0
      INCLUDE = 1
      
  
  
  class SegIEPoly():
      def __init__(self, type=None, pts=None, **kwargs):
          self.type = type
          
          if pts is None:
              pts = np.empty( (0,2), dtype=np.float32 )
          else:
              pts = np.float32(pts)
          self.pts = pts
          self.n_max = self.n = len(pts)
  
      def dump(self):
          return {'type': int(self.type),
                  'pts' : self.get_pts(),
                 }
      
      def identical(self, b):
          if self.n != b.n:
              return False            
          return (self.pts[0:self.n] == b.pts[0:b.n]).all()        
          
      def get_type(self):
          return self.type
  
      def add_pt(self, x, y):
          self.pts = np.append(self.pts[0:self.n], [ ( float(x), float(y) ) ], axis=0).astype(np.float32)
          self.n_max = self.n = self.n + 1
  
      def undo(self):
          self.n = max(0, self.n-1)
          return self.n
  
      def redo(self):
          self.n = min(len(self.pts), self.n+1)
          return self.n
  
      def redo_clip(self):
          self.pts = self.pts[0:self.n]
          self.n_max = self.n
  
      def insert_pt(self, n, pt):
          if n < 0 or n > self.n:
              raise ValueError("insert_pt out of range")
          self.pts = np.concatenate( (self.pts[0:n], pt[None,...].astype(np.float32), self.pts[n:]), axis=0)
          self.n_max = self.n = self.n+1
          
      def remove_pt(self, n):
          if n < 0 or n >= self.n:
              raise ValueError("remove_pt out of range")
          self.pts = np.concatenate( (self.pts[0:n], self.pts[n+1:]), axis=0)
          self.n_max = self.n = self.n-1
          
      def get_last_point(self):
          return self.pts[self.n-1].copy()
  
      def get_pts(self):
          return self.pts[0:self.n].copy()
          
      def get_pts_count(self):
          return self.n
  
      def set_point(self, id, pt):
          self.pts[id] = pt
          
      def set_points(self, pts):
          self.pts = np.array(pts)
          self.n_max = self.n = len(pts)
          
      def mult_points(self, val):
          self.pts *= val       
      
          
  
  class SegIEPolys():
      def __init__(self):
          self.polys = []
  
      def identical(self, b):
          polys_len = len(self.polys)
          o_polys_len = len(b.polys)
          if polys_len != o_polys_len:
              return False
          
          return all ([ a_poly.identical(b_poly) for a_poly, b_poly in zip(self.polys, b.polys) ])
          
      def add_poly(self, ie_poly_type):       
          poly = SegIEPoly(ie_poly_type)
          self.polys.append (poly)
          return poly
  
      def remove_poly(self, poly):
          if poly in self.polys:
              self.polys.remove(poly)
  
      def has_polys(self):
          return len(self.polys) != 0
          
      def get_poly(self, id):
          return self.polys[id]
  
      def get_polys(self):
          return self.polys
          
      def get_pts_count(self):
          return sum([poly.get_pts_count() for poly in self.polys])
          
      def sort(self):
          poly_by_type = { SegIEPolyType.EXCLUDE : [], SegIEPolyType.INCLUDE : [] }
  
          for poly in self.polys:
              poly_by_type[poly.type].append(poly)
              
          self.polys = poly_by_type[SegIEPolyType.INCLUDE] + poly_by_type[SegIEPolyType.EXCLUDE]
  
      def __iter__(self):
          for poly in self.polys:
              yield poly
  
      def overlay_mask(self, mask):
          h,w,c = mask.shape
          white = (1,)*c
          black = (0,)*c
          for poly in self.polys:
              pts = poly.get_pts().astype(np.int32)
              if len(pts) != 0:
                  cv2.fillPoly(mask, [pts], white if poly.type == SegIEPolyType.INCLUDE else black )
  
      def dump(self):
          return {'polys' : [ poly.dump() for poly in self.polys ] }
      
      def mult_points(self, val):
          for poly in self.polys:
              poly.mult_points(val)
          
      @staticmethod
      def load(data=None):
          ie_polys = SegIEPolys()
          if data is not None:
              if isinstance(data, list):
                  # Backward comp
                  ie_polys.polys = [ SegIEPoly(type=type, pts=pts) for (type, pts) in data ]
              elif isinstance(data, dict):                
                  ie_polys.polys = [ SegIEPoly(**poly_cfg) for poly_cfg in data['polys'] ]
                  
          ie_polys.sort()   
                  
          return ie_polys