1 Commits
Author SHA1 Message Date
aravind 79f7e55440 Added oneformer mask 2024-02-01 12:47:01 +05:30
4 changed files with 82 additions and 1276 deletions
-1
View File
@@ -5,6 +5,5 @@ venv
.DS_Store
checkpoints/
checkpoint/
.env
dwpose/keypoints/
+82 -223
View File
@@ -7,6 +7,10 @@ import comfy.model_management as model_management
import folder_paths
from PIL import Image, ImageOps
from transformers import AutoProcessor
from transformers import OneFormerForUniversalSegmentation
from transformers import OneFormerProcessor
def from_torch_image(image):
image = image.squeeze().cpu().numpy() * 255.0
@@ -53,6 +57,30 @@ def do_work_with_mask(image, mask, external_file_name):
return image
def get_segmentation_of_person(image):
NAME_MODEL_ONEFORMER = 'shi-labs/oneformer_ade20k_dinat_large'
processor = OneFormerProcessor.from_pretrained(NAME_MODEL_ONEFORMER)
model = OneFormerForUniversalSegmentation.from_pretrained(
NAME_MODEL_ONEFORMER)
width = image.shape[1]
height = image.shape[0]
inputs = processor(image, ["semantic"], return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
predicted_semantic_map = processor.post_process_semantic_segmentation(
outputs, target_sizes=[(height, width)])[0]
img = predicted_semantic_map.detach().cpu().numpy()
img = (img == 12)
img = img.astype(np.uint8) * 255
return img
class TRI3DATRParseBatch:
def __init__(self):
@@ -1150,15 +1178,15 @@ class TRI3DPoseAdaption:
"required": {
"input_pose_json_file": ("STRING",{"default" : "dwpose/keypoints/input.json"}),
"ref_pose_json_file": ("STRING",{"default" : "dwpose/keypoints/ref-pose.json"}),
"image_angle": (["front", "back","back_fixed_kid","back_fixed","back_fixed_left","back_fixed_right"], {"default": "front"}),
"image_angle": (["front", "back","back_fixed","back_fixed_left","back_fixed_right"], {"default": "front"}),
"rotation_threshold": ("FLOAT", {
"default": 5.0,
"min": 0.0,
"max": 15.0,
"step": 0.01
}),
"garment_category":(["no_sleeve_garment", "half_sleeve_garment", "full_sleeve_garment", \
"shorts", "trouser"], {"default": "no_sleeve_garment"})
"garment_category":([" no_sleeve_garment", "half_sleeve_garment", "full_sleeve_garment", \
"shorts", "trouser"], {"default": " no_sleeve_garment"})
}
@@ -1170,9 +1198,6 @@ class TRI3DPoseAdaption:
def main(self, input_pose_json_file, ref_pose_json_file, image_angle,
rotation_threshold, garment_category):
print(image_angle, "image_angle")
print(garment_category, "garment_category")
from .dwpose import comfy_utils
if image_angle == "front":
@@ -1386,7 +1411,6 @@ class TRI3DPoseAdaption:
back_pose_dir = pathlib.Path().resolve(
) / 'custom_nodes/tri3d-comfyui-nodes/samples/back_poses/'
back_pose_dictionary = {
'back_fixed_kid': 'backpose_kid.json',
'back_fixed': 'backpose.json',
'back_fixed_left': 'left_backpose.json',
'back_fixed_right': 'right_backpose.json'
@@ -1785,124 +1809,6 @@ class FloatToImage:
return image
class TRI3D_recolor_LAB_manual:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"image_input": ("IMAGE", ),
"mask_input": ("IMAGE", ),
"factor_mean": ("FLOAT", {
"default": 1.0,
"min": 0.0,
"max": 10.0,
"step": 0.01
}),
"factor_sigma": ("FLOAT", {
"default": 1.0,
"min": 0.0,
"max": 10.0,
"step": 0.01
}),
}
}
RETURN_TYPES = ("IMAGE", )
FUNCTION = "recolor"
CATEGORY = "TRI3D"
def recolor(self, image_input, mask_input, factor_mean, factor_sigma):
def get_mu_sigma(array_input, mask_input):
import numpy as np
import math
array_input = array_input.astype(dtype=np.float32).flatten()
mask_input = mask_input.flatten()
sum = np.sum(mask_input)
mean = np.sum(array_input * mask_input) / sum
array_input -= mean
array_input *= mask_input
sigma = math.sqrt(np.sum(np.square(array_input)) / sum)
return mean, sigma
def do_recolor(image, mask, mean, sigma):
import cv2
import numpy as np
import math
image_original = image.copy()
mask = (mask > 127).astype(dtype=np.uint8)
sum = np.sum(mask.flatten())
for i in range(3):
image[:, :, i] *= mask
image = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
image = image.astype(dtype=np.float32)
for i in range(1):
mu_1, sigma_1 = get_mu_sigma(array_input=image[:, :, i],
mask_input=mask)
image[:, :, i] = (((image[:, :, i] - mu_1) / sigma_1) *
(sigma * sigma_1)) + (mean * mu_1)
image = np.clip(image, 0, 255)
image = image.astype(dtype=np.uint8)
image = cv2.cvtColor(image, cv2.COLOR_LAB2BGR)
for i in range(3):
image_original[:, :,
i] = (image_original[:, :, i] *
(1 - mask)) + (image[:, :, i] * mask)
return image_original
def from_torch_image(image):
image = image.squeeze().cpu().numpy() * 255.0
image = np.clip(image, 0, 255).astype(np.uint8)
return image
def to_torch_image(image):
import numpy as np
import torch
image = image.astype(dtype=np.float32)
image /= 255.0
image = torch.from_numpy(image)[
None,
]
image = image.unsqueeze(0)
return image
image = from_torch_image(image=image_input)
mask = from_torch_image(image=mask_input)[:, :, 0]
image_output = do_recolor(image=image,
mask=mask,
mean=factor_mean,
sigma=factor_sigma)
image_output = to_torch_image(image=image_output)
return image_output
class TRI3D_recolor_LAB:
@classmethod
@@ -2381,7 +2287,7 @@ class TRI3D_clipdrop_bgremove_api:
return output,
class TRI3DAdjustNeck:
class TRI3D_oneformer_person_mask:
def __init__(self):
pass
@@ -2390,88 +2296,25 @@ class TRI3DAdjustNeck:
def INPUT_TYPES(s):
return {
"required": {
"posemap_json_file_path": ("STRING", {
"default":
"dwpose/keypoints/input.json"
}),
# "age_group" :(["10-12 yrs"],{"default":"10-12 yrs"}),
"neck_shoulder_ratio": ("FLOAT", {
"default": 0.7,
"min": 0.0,
"max": 1.0,
"step": 0.01
}),
"save_json_file_path": ("STRING", {
"default":
"dwpose/keypoints/output.json"
})
"image": ("IMAGE", ),
},
}
FUNCTION = "run"
RETURN_TYPES = ("IMAGE", "STRING")
CATEGORY = "TRI3D"
RETURN_TYPES = ("IMAGE", )
CATEGORY = "tri3d"
def run(self, posemap_json_file_path, neck_shoulder_ratio,
save_json_file_path):
from .dwpose import comfy_utils
def run(self, image):
image = from_torch_image(image)
# age_to_ratio = {"10-12 yrs": 0.65}
image_final = np.zeros((image.shape[0], image.shape[1], 4),
dtype=np.uint8)
input_pose = json.load(open(posemap_json_file_path))
input_height = input_pose['height']
input_width = input_pose['width']
input_keypoints = input_pose['keypoints']
ref_x1, ref_y1 = input_keypoints[2] #left shoulder
ref_x2, ref_y2 = input_keypoints[5] #right_shoulder
x1, y1 = input_keypoints[1] #neck
x2, y2 = input_keypoints[0] #nose
prev_nose = input_keypoints[0]
ref_len = np.linalg.norm(
np.array([ref_x1, ref_y1]) - np.array(
[ref_x2, ref_y2])) #ref body part length i.e. shoulder length
targ_len = np.linalg.norm(np.array([x1, y1]) - np.array([
x2, y2
])) #targ body part length i.e. neck length - neck to nose length
input_targ_ref_len = targ_len / ref_len #neck to shoulder ratio
print("neck to shoulder ratio found", input_targ_ref_len)
print("neck to shoulder ratio target", neck_shoulder_ratio)
#scale the coords
# scale = age_to_ratio[age_group] / input_targ_ref_len
scale = neck_shoulder_ratio / input_targ_ref_len
x2_scaled = x1 + (x2 - x1) * scale
y2_scaled = y1 + (y2 - y1) * scale
input_keypoints[0] = [x2_scaled, y2_scaled]
#changing face points to w.r.t to new nose point after rotation
input_keypoints[14:18] = comfy_utils.move(prev_nose,
input_keypoints[0],
input_keypoints[14:18])
canvas = np.zeros(shape=(input_height, input_width, 3), dtype=np.uint8)
canvas = comfy_utils.draw_bodypose(canvas, input_keypoints)
canvas = comfy_utils.draw_handpose(
canvas, input_keypoints[88:109]) #right hand
canvas = comfy_utils.draw_handpose(canvas,
input_keypoints[109:]) #left hand
canvas = torch.from_numpy(canvas.astype(np.float32) / 255.0)[
None,
]
output_pose = {
"height": input_height,
"width": input_width,
"keypoints": input_keypoints
}
cur_file_dir = os.path.dirname(os.path.realpath(__file__))
save_json_file_path = os.path.join(cur_file_dir, save_json_file_path)
json.dump(output_pose, open(save_json_file_path, 'w'))
return (canvas, save_json_file_path)
image_final[:, :, 0:3] = image
image = get_segmentation_of_person(image)
image_final[:, :, 3] = image[:, :]
image = to_torch_image(image_final)
return image
# A dictionary that contains all nodes you want to export with their names
@@ -2496,34 +2339,50 @@ NODE_CLASS_MAPPINGS = {
"tri3d-image-mask-2-box": TRI3D_image_mask_2_box,
"tri3d-image-mask-box-2-image": TRI3D_image_mask_box_2_image,
"tri3d-clipdrop-bgremove-api": TRI3D_clipdrop_bgremove_api,
"tri3d-adjust-neck": TRI3DAdjustNeck,
"tri3d_recolor_lab_manual": TRI3D_recolor_LAB_manual,
"tri3d-oneformer-person-mask": TRI3D_oneformer_person_mask,
}
VERSION = "2.4.2"
VERSION = "2.3"
# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
"tri3d-atr-parse-batch": "ATR Parse Batch" + " v" + VERSION,
'tri3d-extract-parts-batch': 'Extract Parts Batch' + " v" + VERSION,
'tri3d-extract-parts-batch2': 'Extract Parts Batch 2' + " v" + VERSION,
"tri3d-position-parts-batch": "Position Parts Batch" + " v" + VERSION,
"tri3d-swap-pixels": "Swap Pixels by Mask" + " v" + VERSION,
"tri3d-atr-parse-batch":
"ATR Parse Batch" + " v" + VERSION,
'tri3d-extract-parts-batch':
'Extract Parts Batch' + " v" + VERSION,
'tri3d-extract-parts-batch2':
'Extract Parts Batch 2' + " v" + VERSION,
"tri3d-position-parts-batch":
"Position Parts Batch" + " v" + VERSION,
"tri3d-swap-pixels":
"Swap Pixels by Mask" + " v" + VERSION,
"tri3d-skin-feathered-padded-mask":
"Skin Feathered Padded Mask" + " v" + VERSION,
"tri3d-interaction-canny":
"Garment Skin Interaction Canny" + " v" + VERSION,
"tri3d-dwpose": "DWPose" + " v" + VERSION,
"tri3d-pose-to-image": "Pose to Image" + " v" + VERSION,
"tri3d-pose-adaption": "Pose Adaption" + " v" + VERSION,
"tri3d-load-pose-json": "Load Pose Json" + " v" + VERSION,
"tri3d-face-recognise": "Recognise face" + " v" + VERSION,
"tri3d-float-to-image": "Render float" + " v" + VERSION,
"tri3d-recolor-mask": "Recolor mask HSV space" + " v" + VERSION,
"tri3d-recolor-mask-LAB_space": "Recolor mask LAB space" + " v" + VERSION,
"tri3d-recolor-mask-RGB_space": "Recolor mask RGB space" + " v" + VERSION,
"tri3d--image-mask-2-box": "Extract box from image" + " v" + VERSION,
"tri3d-image-mask-box-2-image": "Stitch box to image" + " v" + VERSION,
"tri3d-clipdrop-bgremove-api": "RemBG ClipDrop" + " v" + VERSION,
"tri3d-adjust-neck": "Adjust Neck" + " v" + VERSION,
"tri3d_recolor_lab_manual": "Adjust color manually" + " v" + VERSION,
"tri3d-dwpose":
"DWPose" + " v" + VERSION,
"tri3d-pose-to-image":
"Pose to Image" + " v" + VERSION,
"tri3d-pose-adaption":
"Pose Adaption" + " v" + VERSION,
"tri3d-load-pose-json":
"Load Pose Json" + " v" + VERSION,
"tri3d-face-recognise":
"Recognise face" + " v" + VERSION,
"tri3d-float-to-image":
"Render float" + " v" + VERSION,
"tri3d-recolor-mask":
"Recolor mask HSV space" + " v" + VERSION,
"tri3d-recolor-mask-LAB_space":
"Recolor mask LAB space" + " v" + VERSION,
"tri3d-recolor-mask-RGB_space":
"Recolor mask RGB space" + " v" + VERSION,
"tri3d--image-mask-2-box":
"Extract box from image" + " v" + VERSION,
"tri3d-image-mask-box-2-image":
"Stitch box to image" + " v" + VERSION,
"tri3d-clipdrop-bgremove-api":
"RemBG ClipDrop" + " v" + VERSION,
"tri3d-oneformer-person-mask":
"Get person mask using oneformer" + " v" + VERSION,
}
-526
View File
@@ -1,526 +0,0 @@
{
"height": 512,
"width": 512,
"keypoints": [
[
-1,
-1
],
[
252,
130
],
[
302,
130
],
[
316,
213
],
[
323,
284
],
[
202,
130
],
[
191,
211
],
[
187,
282
],
[
282,
278
],
[
280,
380
],
[
280,
484
],
[
222,
278
],
[
226,
382
],
[
222,
480
],
[
-1,
-1
],
[
-1,
-1
],
[
286,
72
],
[
227,
73
],
[
257,
59
],
[
257,
63
],
[
257,
65
],
[
253,
67
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
279,
79
],
[
-2,
-2
],
[
243,
67
],
[
253,
71
],
[
257,
69
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
275,
61
],
[
279,
59
],
[
255,
55
],
[
-2,
-2
],
[
-2,
-2
],
[
255,
59
],
[
255,
59
],
[
255,
59
],
[
255,
57
],
[
255,
57
],
[
255,
59
],
[
255,
59
],
[
255,
61
],
[
255,
63
],
[
277,
67
],
[
275,
67
],
[
277,
67
],
[
277,
67
],
[
277,
67
],
[
255,
59
],
[
255,
57
],
[
257,
59
],
[
257,
61
],
[
257,
59
],
[
255,
59
],
[
255,
59
],
[
257,
59
],
[
257,
59
],
[
-2,
-2
],
[
255,
59
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
279,
71
],
[
243,
67
],
[
279,
73
],
[
277,
71
],
[
279,
71
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
-2,
-2
],
[
279,
69
],
[
279,
71
],
[
277,
71
],
[
279,
71
],
[
277,
69
],
[
-2,
-2
],
[
250,
56
],
[
263,
63
],
[
186.65255255416977,
285
],
[
196.65255255416977,
291
],
[
199.65255255416977,
298
],
[
201.65255255416977,
306
],
[
201.65255255416977,
310
],
[
194.65255255416977,
308
],
[
196.65255255416977,
312
],
[
199.65255255416977,
316
],
[
203.65255255416977,
318
],
[
190.65255255416977,
308
],
[
192.65255255416977,
314
],
[
194.65255255416977,
318
],
[
201.65255255416977,
320
],
[
186.65255255416977,
308
],
[
186.65255255416977,
314
],
[
190.65255255416977,
318
],
[
196.65255255416977,
320
],
[
182.65255255416977,
308
],
[
184.65255255416977,
312
],
[
186.65255255416977,
316
],
[
190.65255255416977,
318
],
[
323,
287
],
[
313,
298
],
[
310,
305
],
[
310,
312
],
[
310,
317
],
[
320,
315
],
[
317,
319
],
[
313,
324
],
[
308,
326
],
[
322,
315
],
[
320,
322
],
[
315,
326
],
[
308,
326
],
[
324,
317
],
[
320,
322
],
[
315,
326
],
[
310,
326
],
[
327,
317
],
[
324,
322
],
[
320,
324
],
[
315,
326
]
]
}
-526
View File
@@ -1,526 +0,0 @@
{
"height": 512,
"width": 512,
"keypoints": [
[
-1,
-1
],
[
254,
118
],
[
292,
118
],
[
306,
194
],
[
310,
262
],
[
216,
118
],
[
200,
193
],
[
195,
260
],
[
278,
251
],
[
286,
357
],
[
285.99326159010275,
462.45102588957826
],
[
228,
254
],
[
230,
353
],
[
231.99326159010275,
458.45102588957826
],
[
-1,
-1
],
[
-1,
-1
],
[
277,
62
],
[
241,
62
],
[
234,
53
],
[
234,
60
],
[
234,
64
],
[
236,
71
],
[
236,
75
],
[
241,
77
],
[
243,
82
],
[
247,
84
],
[
254,
84
],
[
258,
84
],
[
263,
84
],
[
267,
82
],
[
272,
80
],
[
274,
75
],
[
276,
71
],
[
278,
66
],
[
281,
62
],
[
243,
44
],
[
245,
42
],
[
247,
42
],
[
252,
42
],
[
256,
42
],
[
265,
44
],
[
269,
44
],
[
272,
46
],
[
276,
49
],
[
278,
51
],
[
261,
51
],
[
258,
53
],
[
258,
55
],
[
258,
55
],
[
254,
62
],
[
254,
62
],
[
258,
62
],
[
261,
62
],
[
263,
62
],
[
245,
49
],
[
247,
49
],
[
252,
49
],
[
254,
49
],
[
249,
51
],
[
247,
49
],
[
265,
51
],
[
269,
51
],
[
272,
51
],
[
274,
55
],
[
272,
55
],
[
267,
53
],
[
247,
69
],
[
252,
66
],
[
254,
66
],
[
256,
66
],
[
258,
66
],
[
263,
69
],
[
265,
71
],
[
263,
73
],
[
258,
73
],
[
256,
73
],
[
252,
73
],
[
249,
71
],
[
247,
69
],
[
252,
66
],
[
256,
66
],
[
261,
69
],
[
265,
71
],
[
261,
73
],
[
256,
71
],
[
252,
71
],
[
249,
49
],
[
269,
51
],
[
195,
262
],
[
206,
270
],
[
210,
276
],
[
213,
280
],
[
216,
285
],
[
201,
287
],
[
208,
294
],
[
213,
296
],
[
216,
299
],
[
197,
289
],
[
206,
296
],
[
208,
299
],
[
213,
299
],
[
195,
289
],
[
201,
294
],
[
208,
296
],
[
208,
299
],
[
195,
289
],
[
197,
294
],
[
201,
296
],
[
206,
296
],
[
310,
265
],
[
304,
270
],
[
302,
274
],
[
299,
279
],
[
297,
284
],
[
308,
287
],
[
307,
292
],
[
304,
299
],
[
299,
301
],
[
309,
287
],
[
308,
292
],
[
307,
299
],
[
303,
301
],
[
310,
287
],
[
309,
292
],
[
308,
298
],
[
307,
302
],
[
313,
287
],
[
310,
291
],
[
309,
296
],
[
308,
299
]
]
}