Podtekatel commited on
Commit
f02b231
β€’
1 Parent(s): de7b3c0

Update infer pipeline

Browse files
Files changed (1) hide show
  1. inference/model_pipeline.py +8 -3
inference/model_pipeline.py CHANGED
@@ -9,12 +9,13 @@ from .face_detector import FaceDetector
9
 
10
 
11
  class VSNetModelPipeline:
12
- def __init__(self, model, face_detector: FaceDetector, background_resize=720, no_detected_resize=256):
13
  self.background_resize = background_resize
14
  self.no_detected_resize = no_detected_resize
15
  self.model = model
16
  self.face_detector = face_detector
17
  self.mask = self.create_circular_mask(face_detector.target_size, face_detector.target_size)
 
18
 
19
  @staticmethod
20
  def create_circular_mask(h, w, power=None, clipping_coef=0.85):
@@ -68,8 +69,12 @@ class VSNetModelPipeline:
68
  W, H = x2 - x1, y2 - y1
69
  result_face = cv2.resize(face, (W, H), interpolation=cv2.INTER_LINEAR)
70
  face_mask = cv2.resize(self.mask, (W, H), interpolation=cv2.INTER_LINEAR)
71
- input_face = full_image[y1: y2, x1: x2]
72
- full_image[y1: y2, x1: x2] = (result_face * face_mask + input_face * (1 - face_mask)).astype(np.uint8)
 
 
 
 
73
  return full_image
74
 
75
  def __call__(self, img):
 
9
 
10
 
11
  class VSNetModelPipeline:
12
+ def __init__(self, model, face_detector: FaceDetector, background_resize=720, no_detected_resize=256, use_cloning=True):
13
  self.background_resize = background_resize
14
  self.no_detected_resize = no_detected_resize
15
  self.model = model
16
  self.face_detector = face_detector
17
  self.mask = self.create_circular_mask(face_detector.target_size, face_detector.target_size)
18
+ self.use_cloning = use_cloning
19
 
20
  @staticmethod
21
  def create_circular_mask(h, w, power=None, clipping_coef=0.85):
 
69
  W, H = x2 - x1, y2 - y1
70
  result_face = cv2.resize(face, (W, H), interpolation=cv2.INTER_LINEAR)
71
  face_mask = cv2.resize(self.mask, (W, H), interpolation=cv2.INTER_LINEAR)
72
+ if self.use_cloning:
73
+ center = round((x2 + x1) / 2), round((y2 + y1) / 2)
74
+ full_image = cv2.seamlessClone(result_face, full_image, (face_mask > 0.0).astype(np.uint8) * 255, center, cv2.NORMAL_CLONE)
75
+ else:
76
+ input_face = full_image[y1: y2, x1: x2]
77
+ full_image[y1: y2, x1: x2] = (result_face * face_mask + input_face * (1 - face_mask)).astype(np.uint8)
78
  return full_image
79
 
80
  def __call__(self, img):