Update app.py

app.py

@@ -50,7 +50,7 @@ PART_GROUPS = {
     "mouth": {"indices": [24, 25, 26, 27],          "use_box": True},
     "browL": {"indices": [8, 9, 10],                "use_box": True},
     "browR": {"indices": [5, 6, 7],                 "use_box": True},
-    "face":  {"indices": [0, 1, 2, 3, 4],           "use_box": False},  # 顎だけなのでボックス不向き
+    # face はSAMを使わず幾何学的に生成（後処理で追加）
 }
 
 
@@ -144,6 +144,53 @@ def segment_part(image: Image.Image, image_embeddings, part_points: List[List[fl
     return (best_mask > 0).astype(np.uint8) * 255
 
 
+def compute_face_mask(points: List[dict], image_size: Tuple[int, int]) -> np.ndarray:
+    """顎(0-4)の輪郭を上側で円弧で閉じた顔マスクを生成"""
+    from PIL import ImageDraw
+    import math
+
+    w, h = image_size
+
+    # 顎の5点
+    jaw = [(points[i]["x"], points[i]["y"]) for i in range(5) if i < len(points)]
+    if len(jaw) < 3:
+        return np.zeros((h, w), dtype=np.uint8)
+
+    # 眉の上端 → 額の位置
+    brow_ys = [points[i]["y"] for i in range(5, 11) if i < len(points)]
+    forehead_y = min(brow_ys) if brow_ys else jaw[0][1]
+    # 額にマージンを追加
+    face_width = abs(jaw[0][0] - jaw[-1][0])
+    top_y = forehead_y - face_width * 0.2
+
+    # 左端(jaw[0])と右端(jaw[-1])を円弧で繋ぐ
+    # 上部の制御点を生成（半円弧を点列で近似）
+    left = jaw[0]   # 左の顎端
+    right = jaw[-1]  # 右の顎端
+    cx = (left[0] + right[0]) / 2
+    rx = abs(right[0] - left[0]) / 2 * 1.1  # 少し広め
+    ry = max(left[1], right[1]) - top_y
+
+    # 右端→上→左端の円弧（反時計回り）
+    arc_points = []
+    n_arc = 16
+    for i in range(n_arc + 1):
+        angle = math.pi * i / n_arc  # 0 → π（右→上→左）
+        ax = cx + rx * math.cos(angle)
+        ay = min(left[1], right[1]) - ry * math.sin(angle)
+        arc_points.append((ax, ay))
+
+    # 完全なポリゴン: 顎(0→4) + 円弧(右端→上→左端)
+    polygon = list(jaw) + arc_points
+
+    # ポリゴンを塗りつぶしてマスク生成
+    mask_img = Image.new("L", (w, h), 0)
+    draw = ImageDraw.Draw(mask_img)
+    draw.polygon(polygon, fill=255)
+
+    return np.array(mask_img)
+
+
 def estimate_hair_points(points: List[dict], image_size: Tuple[int, int]) -> List[List[float]]:
     """ランドマークから髪の推定ポイントを生成（頭上 + 左右 + サイド）"""
     w, h = image_size
@@ -211,19 +258,10 @@ def run_with_landmarks(image: Image.Image, landmarks_json: str) -> Tuple[List[st
     status_parts = []
     all_masks = []  # 全パーツマスクを蓄積（残り領域計算用）
 
-    # 髪のポイントを先に計算（face のネガティブプロンプトに使う）
-    hair_points = estimate_hair_points(points, image.size)
-
-    # 各パーツをセグメント
+    # 各パーツをセグメント（目・口・眉のみ — SAMが得意な小パーツ）
     for part_name, config in PART_GROUPS.items():
         part_points = get_all_points(points, config["indices"])
-
-        # face: 髪の推定ポイントをネガティブに渡して髪領域を除外
-        neg = None
-        if part_name == "face" and hair_points:
-            neg = hair_points[:3]  # 頭頂3点を背景ポイントに
-
-        mask_arr = segment_part(image, image_embeddings, part_points, original_size, reshaped_size, use_box=config.get("use_box", True), negative_points=neg)
+        mask_arr = segment_part(image, image_embeddings, part_points, original_size, reshaped_size, use_box=config.get("use_box", True))
 
         if mask_arr is not None:
             mask_img = Image.fromarray(mask_arr, "L")
@@ -236,54 +274,33 @@ def run_with_landmarks(image: Image.Image, landmarks_json: str) -> Tuple[List[st
             status_parts.append(part_name)
             all_masks.append(mask_arr)
 
-    # 髪のセグメント — 顔の中心をネガティブポイントにして顔領域を除外
-    if hair_points:
-        jaw_points = get_all_points(points, [0, 1, 2, 3, 4])
-        face_center = None
-        if jaw_points:
-            face_center = [[
-                sum(p[0] for p in jaw_points) / len(jaw_points),
-                sum(p[1] for p in jaw_points) / len(jaw_points),
-            ]]
-        hair_mask = segment_part(image, image_embeddings, hair_points, original_size, reshaped_size, use_box=False, negative_points=face_center)
-        if hair_mask is not None:
-            hair_img = Image.fromarray(hair_mask, "L")
-            if hair_img.size != image.size:
-                hair_img = hair_img.resize(image.size)
-                hair_mask = np.array(hair_img)
-            path = os.path.join(out_dir, "hair.png")
-            hair_img.save(path)
-            mask_files.append(path)
-            status_parts.append("hair")
-            all_masks.append(hair_mask)
-
-    # ===== 後処理: マスク同士の重なりを除去 =====
-    # 優先順位: 小パーツ（目・眉・口） > 顔 > 髪
-    # 小パー��領域を顔・髪から引き算し、顔領域を髪から引き算
+    # ===== 後処理: 顔を幾何学的に追加 + 重なり除去 =====
     part_masks = dict(zip(status_parts, all_masks))
     small_parts = ["eyeL", "eyeR", "mouth", "browL", "browR"]
+    h_img, w_img = all_masks[0].shape if all_masks else (image.size[1], image.size[0])
+
+    # 幾何学的な顔マスク（楕円）を生成、小パーツを除外
+    face_ellipse = compute_face_mask(points, image.size)
+    if face_ellipse.shape != (h_img, w_img):
+        face_ellipse = np.array(Image.fromarray(face_ellipse, "L").resize((w_img, h_img)))
 
     # 小パーツの和集合
-    small_union = np.zeros_like(all_masks[0]) if all_masks else None
+    small_union = np.zeros((h_img, w_img), dtype=np.uint8)
     for sp in small_parts:
         if sp in part_masks:
             small_union = np.maximum(small_union, part_masks[sp])
 
-    # 顔から小パーツを引く
-    if "face" in part_masks and small_union is not None:
-        part_masks["face"] = np.where(small_union > 0, 0, part_masks["face"]).astype(np.uint8)
-
-    # 髪から顔+小パーツを引く
-    if "hair" in part_masks:
-        exclude = small_union.copy() if small_union is not None else np.zeros_like(all_masks[0])
-        if "face" in part_masks:
-            exclude = np.maximum(exclude, part_masks["face"])
-        part_masks["hair"] = np.where(exclude > 0, 0, part_masks["hair"]).astype(np.uint8)
+    # 顔 = 楕円マスク − 小パーツ
+    face_mask = np.where(small_union > 0, 0, face_ellipse).astype(np.uint8)
+    part_masks["face"] = face_mask
 
     # 修正後のマスクを再保存
     mask_files = []
     all_masks = []
-    for part_name in status_parts:
+    status_parts_clean = [p for p in status_parts]
+    if "face" not in status_parts_clean:
+        status_parts_clean.append("face")
+    for part_name in status_parts_clean:
         arr = part_masks[part_name]
         mask_img = Image.fromarray(arr, "L")
         path = os.path.join(out_dir, f"{part_name}.png")
@@ -291,20 +308,65 @@ def run_with_landmarks(image: Image.Image, landmarks_json: str) -> Tuple[List[st
         mask_files.append(path)
         all_masks.append(arr)
 
-    # 残り領域マスク（body）: どのパーツにも属さない部分
-    if all_masks:
-        union = np.zeros(all_masks[0].shape, dtype=np.uint8)
-        for m in all_masks:
-            union = np.maximum(union, m)
-        remainder = np.where(union > 0, 0, 255).astype(np.uint8)
-        # 残りが意味のあるサイズ（全体の1%以上）なら保存
-        if np.sum(remainder > 0) > remainder.size * 0.01:
+    # ===== 残り領域を「髪」と「体」に分割 =====
+    union = np.zeros((h_img, w_img), dtype=np.uint8)
+    for m in all_masks:
+        union = np.maximum(union, m)
+    remainder = np.where(union > 0, 0, 255).astype(np.uint8)
+
+    if np.sum(remainder > 0) > remainder.size * 0.005:
+        # 顔の中心Y座標を基準に上下分割
+        jaw_ys = [points[i]["y"] for i in range(5) if i < len(points)]
+        jaw_xs = [points[i]["x"] for i in range(5) if i < len(points)]
+        if jaw_ys:
+            face_center_y = sum(jaw_ys) / len(jaw_ys)
+            face_center_x = sum(jaw_xs) / len(jaw_xs)
+            face_w = (max(jaw_xs) - min(jaw_xs)) if len(jaw_xs) > 1 else w_img * 0.4
+
+            # 髪 = 残りのうち、顔の上部 or 顔の横（顔中心Yより上 + 顔幅の外側で上半分）
+            hair_mask = np.zeros_like(remainder)
+            for y in range(h_img):
+                for x in range(w_img):
+                    if remainder[y, x] == 0:
+                        continue
+                    # 顔中心より上 → 髪
+                    if y < face_center_y:
+                        hair_mask[y, x] = 255
+                    # 顔の横で上半分 → サイドの髪
+                    elif y < face_center_y + face_w * 0.5:
+                        dist_from_center = abs(x - face_center_x)
+                        if dist_from_center > face_w * 0.35:
+                            hair_mask[y, x] = 255
+
+            # 体 = 残り − 髪
+            body_mask = np.where(hair_mask > 0, 0, remainder).astype(np.uint8)
+
+            # 髪マスクを保存（意味のあるサイズなら）
+            if np.sum(hair_mask > 0) > remainder.size * 0.005:
+                hair_img = Image.fromarray(hair_mask, "L")
+                path = os.path.join(out_dir, "hair.png")
+                hair_img.save(path)
+                mask_files.append(path)
+                status_parts_clean.append("hair")
+                all_masks.append(hair_mask)
+
+            # 体マスクを保存
+            if np.sum(body_mask > 0) > remainder.size * 0.005:
+                body_img = Image.fromarray(body_mask, "L")
+                path = os.path.join(out_dir, "body.png")
+                body_img.save(path)
+                mask_files.append(path)
+                status_parts_clean.append("body")
+                all_masks.append(body_mask)
+        else:
+            # ランドマークなしの場合は全部bodyに
             remainder_img = Image.fromarray(remainder, "L")
             path = os.path.join(out_dir, "body.png")
             remainder_img.save(path)
             mask_files.append(path)
-            status_parts.append("body")
-            all_masks.append(np.array(remainder_img))
+            status_parts_clean.append("body")
+
+    status_parts = status_parts_clean
 
     # ZIP
     if mask_files: