mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2024-12-25 07:21:13 -08:00
f387179cba
4.2) data_src util faceset resize.bat 5.2) data_dst util faceset resize.bat Resize faceset to match model resolution to reduce CPU load during training. Don’t forget to keep original faceset.
158 lines
4.3 KiB
Python
158 lines
4.3 KiB
Python
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 |