nodes/TRI3D-LC-tri3d-comfyui-nodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

resizing mask to image size in extarct parts and using the same coords to replace extracted part in position node

Browse Source
This commit is contained in:
Ram Deshmukh
2024-05-02 17:53:12 +05:30
parent a7810c118b
commit c335e3dd8c
1 changed files with 31 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+31 -13
View File
@@ -353,7 +353,11 @@ class TRI3DExtractPartsBatch:
RETURN_TYPES = (
"IMAGE",
"IMAGE",
"LIST"
)
RETURN_NAMES = ("extracted_images","extracted_masks","extracted_coords")
FUNCTION = "main"
CATEGORY = "TRI3D"
@@ -443,11 +447,14 @@ class TRI3DExtractPartsBatch:
for i in range(batch_images.shape[0]):
image = batch_images[i]
seg = batch_segs[i]
cv2_image = tensor_to_cv2_img(image)
cv2_secondary = tensor_to_cv2_img(batch_secondaries[i])
cv2_seg = tensor_to_cv2_img(seg)
cv2_secondary = cv2.resize(cv2_secondary, (cv2_image.shape[1], image.shape[0]),
interpolation=cv2.INTER_AREA)
color_code_list = []
################# ATR MAPPING#################
if right_leg:
@@ -530,12 +537,14 @@ class TRI3DExtractPartsBatch:
batch_results = []
batch_secondaries = []
batch_coords = []
for i,img in enumerate(images):
x,y,w,h = images_xywh[i]
img = img[y:y+max_height, x:x+max_width]
img = cv2_img_to_tensor(img)
# print("xywh", x, y, max_width, max_height)
batch_coords.append([x, y, max_width, max_height])
batch_results.append(img.squeeze(0))
for i,sec in enumerate(secondaries):
@@ -550,7 +559,8 @@ class TRI3DExtractPartsBatch:
batch_secondaries = torch.stack(batch_secondaries)
print(batch_results.shape, "batch_results.shape")
return (batch_results, batch_secondaries)
print(batch_coords, "batch_coords")
return (batch_results, batch_secondaries, batch_coords)
class TRI3DPositionPartsBatch:
@@ -563,8 +573,9 @@ class TRI3DPositionPartsBatch:
return {
"required": {
"batch_images": ("IMAGE", ),
"batch_segs": ("IMAGE", ),
# "batch_segs": ("IMAGE", ),
"batch_handimgs": ("IMAGE", ),
"batch_coords":("LIST", ),
"margin": ("INT", {
"default": 15,
"min": 0
@@ -630,7 +641,7 @@ class TRI3DPositionPartsBatch:
FUNCTION = "main"
CATEGORY = "TRI3D"
def main(self, batch_images, batch_segs, batch_handimgs, margin, right_leg,
def main(self, batch_images, batch_handimgs, batch_coords, margin, right_leg,
right_hand, head, hair, left_shoe, bag, background, dress,
left_leg, right_shoe, left_hand, upper_garment, lower_garment,
belt, skirt, hat, sunglasses, scarf):
@@ -698,11 +709,12 @@ class TRI3DPositionPartsBatch:
for i in range(batch_images.shape[0]):
image = batch_images[i]
seg = batch_segs[i]
# seg = batch_segs[i]
handimg = batch_handimgs[i]
positions = batch_coords[i]
cv2_image = tensor_to_cv2_img(image)
cv2_seg = tensor_to_cv2_img(seg)
# cv2_seg = tensor_to_cv2_img(seg)
color_code_list = []
################# ATR MAPPING#################
@@ -743,14 +755,20 @@ class TRI3DPositionPartsBatch:
if scarf:
color_code_list.append([128, 64, 0])
positions = bounded_image_points(cv2_seg, color_code_list,
cv2_image)
# positions = bounded_image_points(cv2_seg, color_code_list,
# cv2_image)
try:
cv2_handimg = tensor_to_cv2_img(handimg)
cv2_handimg = cv2.resize(cv2_handimg,
(positions[2], positions[3]),
interpolation=cv2.INTER_AREA)
# positions = "".join(positions.split(","))
# positions = [int(i) for i in positions]
# print("xywh", positions)
if cv2_handimg.shape[0] != positions[3] or cv2_handimg.shape[1] != positions[2]:
cv2_handimg = cv2.resize(cv2_handimg,
(positions[2], positions[3]),
interpolation=cv2.INTER_AREA)
cv2_handimg = unsharp_mask(cv2_handimg)