ekhatskevich commited on
Commit
3d23955
·
1 Parent(s): 9cda2f8

deal with mask

Browse files
Files changed (2) hide show
  1. app.py +30 -26
  2. modules/ace_plus_ldm.py +4 -0
app.py CHANGED
@@ -53,32 +53,36 @@ def create_face_mask(pil_image):
53
  Create a binary mask (PIL Image) from a PIL image by detecting the face region.
54
  The mask will be white (255) on the detected face area and black (0) elsewhere.
55
  """
56
- # Convert PIL image to a numpy array in RGB format
57
- image_np = np.array(pil_image.convert("RGB"))
58
- # Convert to grayscale for face detection
59
- gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
60
-
61
- # Load the Haar cascade for face detection (make sure opencv data is installed)
62
- cascade_path = cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
63
- face_cascade = cv2.CascadeClassifier(cascade_path)
64
-
65
- # Detect faces in the image
66
- faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
67
-
68
- # Create an empty mask with the same dimensions as the image
69
- mask = np.zeros_like(gray, dtype=np.uint8)
70
-
71
- # For each detected face, draw a white rectangle (or a more refined shape)
72
- for (x, y, w, h) in faces:
73
- # Optionally expand the bounding box slightly
74
- padding = 0.2
75
- x1 = max(0, int(x - w * padding))
76
- y1 = max(0, int(y - h * padding))
77
- x2 = min(gray.shape[1], int(x + w * (1 + padding)))
78
- y2 = min(gray.shape[0], int(y + h * (1 + padding)))
79
- mask[y1:y2, x1:x2] = 255
80
-
81
- return Image.fromarray(mask)
 
 
 
 
82
 
83
  def face_swap_app(target_img, face_img):
84
  if target_img is None or face_img is None:
 
53
  Create a binary mask (PIL Image) from a PIL image by detecting the face region.
54
  The mask will be white (255) on the detected face area and black (0) elsewhere.
55
  """
56
+ try:
57
+ # Convert PIL image to a numpy array in RGB format
58
+ image_np = np.array(pil_image.convert("RGB"))
59
+ # Convert to grayscale for face detection
60
+ gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
61
+
62
+ # Load the Haar cascade for face detection (make sure opencv data is installed)
63
+ cascade_path = cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
64
+ face_cascade = cv2.CascadeClassifier(cascade_path)
65
+
66
+ # Detect faces in the image
67
+ faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5)
68
+
69
+ # Create an empty mask with the same dimensions as the image
70
+ mask = np.zeros_like(gray, dtype=np.uint8)
71
+
72
+ # For each detected face, draw a white rectangle (or a more refined shape)
73
+ for (x, y, w, h) in faces:
74
+ # Optionally expand the bounding box slightly
75
+ padding = 0.2
76
+ x1 = max(0, int(x - w * padding))
77
+ y1 = max(0, int(y - h * padding))
78
+ x2 = min(gray.shape[1], int(x + w * (1 + padding)))
79
+ y2 = min(gray.shape[0], int(y + h * (1 + padding)))
80
+ mask[y1:y2, x1:x2] = 255
81
+
82
+ return Image.fromarray(mask)
83
+ except Exception as e:
84
+ print(f"Error: {e}")
85
+ raise ValueError('A very specific bad thing happened.')
86
 
87
  def face_swap_app(target_img, face_img):
88
  if target_img is None or face_img is None:
modules/ace_plus_ldm.py CHANGED
@@ -100,6 +100,10 @@ class LatentDiffusionACEPlus(LatentDiffusion):
100
  @torch.no_grad()
101
  def encode_first_stage(self, x, **kwargs):
102
  def run_one_image(u):
 
 
 
 
103
  zu = self.first_stage_model.encode(u)
104
  if isinstance(zu, (tuple, list)):
105
  zu = zu[0]
 
100
  @torch.no_grad()
101
  def encode_first_stage(self, x, **kwargs):
102
  def run_one_image(u):
103
+ if u is None:
104
+ print(f"Warning: input is None!")
105
+ else:
106
+ print(f"Input shape: {u.shape}")
107
  zu = self.first_stage_model.encode(u)
108
  if isinstance(zu, (tuple, list)):
109
  zu = zu[0]