Podtekatel commited on
Commit
08befe8
1 Parent(s): 08d07b3

Fix power of mask for inference

Browse files
Files changed (1) hide show
  1. inference/model_pipeline.py +11 -7
inference/model_pipeline.py CHANGED
@@ -14,24 +14,28 @@ class VSNetModelPipeline:
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, power=1 / 4)
18
 
19
  @staticmethod
20
- def create_circular_mask(h, w, center=None, power=None):
21
-
22
- if center is None: # use the middle of the image
23
- center = (int(w / 2), int(h / 2))
24
 
25
  Y, X = np.ogrid[:h, :w]
26
  dist_from_center = np.sqrt((X - center[0]) ** 2 + (Y - center[1]) ** 2)
27
- dist_from_center = np.clip(dist_from_center, a_min=0, a_max=max(h / 2, w / 2))
28
- dist_from_center = 1 - dist_from_center / np.max(dist_from_center)
 
 
 
 
 
29
  if power is not None:
30
  dist_from_center = np.power(dist_from_center, power)
31
  dist_from_center = np.stack([dist_from_center] * 3, axis=2)
32
  # mask = dist_from_center <= radius
33
  return dist_from_center
34
 
 
35
  @staticmethod
36
  def resize_size(image, size=720, always_apply=True):
37
  h, w, c = np.shape(image)
 
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):
21
+ center = (int(w / 2), int(h / 2))
 
 
22
 
23
  Y, X = np.ogrid[:h, :w]
24
  dist_from_center = np.sqrt((X - center[0]) ** 2 + (Y - center[1]) ** 2)
25
+ print(dist_from_center.max(), dist_from_center.min())
26
+ clipping_radius = min((h - center[0]), (w - center[1])) * clipping_coef
27
+ max_size = max((h - center[0]), (w - center[1]))
28
+ dist_from_center[dist_from_center < clipping_radius] = clipping_radius
29
+ dist_from_center[dist_from_center > max_size] = max_size
30
+ max_distance, min_distance = np.max(dist_from_center), np.min(dist_from_center)
31
+ dist_from_center = 1 - (dist_from_center - min_distance) / (max_distance - min_distance)
32
  if power is not None:
33
  dist_from_center = np.power(dist_from_center, power)
34
  dist_from_center = np.stack([dist_from_center] * 3, axis=2)
35
  # mask = dist_from_center <= radius
36
  return dist_from_center
37
 
38
+
39
  @staticmethod
40
  def resize_size(image, size=720, always_apply=True):
41
  h, w, c = np.shape(image)