From b527faff70d4175b7d142b9fbc1f18a3eda3fc05 Mon Sep 17 00:00:00 2001
From: toto <tototianhao@gmail.com>
Date: Wed, 8 Nov 2023 20:38:27 +0800
Subject: [PATCH] bugfix facecrop and segment

---
 comfyui/nodes.py           | 13 +++++++------
 comfyui/utils/img_utils.py | 12 ++++++++++--
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/comfyui/nodes.py b/comfyui/nodes.py
index 0f6b3f6..b46fcae 100644
--- a/comfyui/nodes.py
+++ b/comfyui/nodes.py
@@ -14,8 +14,8 @@ from modelscope.utils.constant import Tasks
 from torch import multiprocessing
 from transformers import pipeline as tpipeline
 
-# import pydevd_pycharm
-# pydevd_pycharm.settrace('49.7.62.197', port=10090, stdoutToServer=True, stderrToServer=True)
+import pydevd_pycharm
+pydevd_pycharm.settrace('49.7.62.197', port=10090, stdoutToServer=True, stderrToServer=True)
 
 from .model_holder import *
 from .utils.img_utils import *
@@ -139,7 +139,7 @@ class FCCropMask:
         inpaint_img = Image.fromarray(cv2.cvtColor(inpaint_img[:, :, ::-1], cv2.COLOR_BGR2RGB))
         mask_large1[cy - cropup:cy + cropbo, cx - crople:cx + cropri] = 1
         mask_large = mask_large * mask_large1
-        return (img_to_tensor(inpaint_img), np_to_mask(mask_large))
+        return (img_to_tensor(inpaint_img), mask_np3_to_mask_tensor(mask_large))
 
 class FCSegment:
     @classmethod
@@ -154,7 +154,7 @@ class FCSegment:
     FUNCTION = "fc_segment"
     CATEGORY = "facechain/model"
 
-    def segment(segmentation_pipeline, img, ksize=0, eyeh=0, ksize1=0, include_neck=False, warp_mask=None, return_human=False):
+    def segment(self, segmentation_pipeline, img, ksize=0, eyeh=0, ksize1=0, include_neck=False, warp_mask=None, return_human=False):
         if True:
             result = segmentation_pipeline(img)
             masks = result['masks']
@@ -185,6 +185,7 @@ class FCSegment:
             mask_neck = np.clip(mask_neck, 0, 1)
             mask_cloth = np.clip(mask_cloth, 0, 1)
             mask_human = np.clip(mask_human, 0, 1)
+            soft_mask = 0
             if np.sum(mask_face) > 0:
                 soft_mask = np.clip(mask_face, 0, 1)
                 if ksize1 > 0:
@@ -218,6 +219,6 @@ class FCSegment:
             return soft_mask
 
     def fc_segment(self, source_image):
-        source_image = img_to_tensor(source_image)
+        source_image = tensor_to_img(source_image)
         mask = self.segment(get_segmentation(), source_image, ksize=0.1)
-        return (img_to_mask(mask),)
+        return (mask_np2_to_mask_tensor(mask),)
diff --git a/comfyui/utils/img_utils.py b/comfyui/utils/img_utils.py
index 98399b8..cf06dfd 100644
--- a/comfyui/utils/img_utils.py
+++ b/comfyui/utils/img_utils.py
@@ -23,11 +23,19 @@ def img_to_mask(input):
     mask_tensor = torch.from_numpy(new_np).permute(2, 0, 1)[0:1, :, :]
     return mask_tensor
 
-def np_to_tensor(input):
+def image_np2_to_mask_tensor(input):
     image = input.astype(np.float32) / 255.0
     tensor = torch.from_numpy(image)[None,]
     return tensor
 
+def mask_np2_to_mask_tensor(input):
+    image = input.astype(np.float32)
+    tensor = torch.from_numpy(image)[None,]
+    return tensor
+def mask_np3_to_mask_tensor(input):
+    image = input.astype(np.float32)
+    tensor = torch.from_numpy(image).permute(2, 0, 1)[0:1, :, :]
+    return tensor
 def tensor_to_img(image):
     image = image[0]
     i = 255. * image.cpu().numpy()
@@ -40,7 +48,7 @@ def tensor_to_np(image):
     result = np.clip(i, 0, 255).astype(np.uint8)
     return result
 
-def np_to_mask(input):
+def image_np_to_mask(input):
     new_np = input.astype(np.float32) / 255.0
     tensor = torch.from_numpy(new_np).permute(2, 0, 1)[0:1, :, :]
     return tensor