xinyun/ComfyUI-KJNodes
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git synced 2026-08-12 09:43:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

BatchCropFromMaskAdvanced: Fixing CenterCrop, handling empty masks.

Browse Source
This commit is contained in:
David Bielejeski 2024-04-03 13:32:12 -05:00
parent b14620f227
commit 4e3fd9ca88
1 changed files with 65 additions and 50 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
nodes.py
View File

@ -1858,8 +1858,13 @@ class BatchCropFromMaskAdvanced:
def calculate_bbox(mask): def calculate_bbox(mask):
non_zero_indices = np.nonzero(np.array(mask)) non_zero_indices = np.nonzero(np.array(mask))
# handle empty masks
min_x, max_x, min_y, max_y = 0, 0, 0, 0
if len(non_zero_indices[1]) > 0 and len(non_zero_indices[0]) > 0:
min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1]) min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1])
min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0]) min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0])
width = max_x - min_x width = max_x - min_x
height = max_y - min_y height = max_y - min_y
bbox_size = max(width, height) bbox_size = max(width, height)
@ -1897,6 +1902,15 @@ class BatchCropFromMaskAdvanced:
for i, (mask, img) in enumerate(zip(masks, original_images)): for i, (mask, img) in enumerate(zip(masks, original_images)):
_mask = tensor2pil(mask)[0] _mask = tensor2pil(mask)[0]
non_zero_indices = np.nonzero(np.array(_mask)) non_zero_indices = np.nonzero(np.array(_mask))
# handle empty masks
if len(non_zero_indices[0]) == 0 or len(non_zero_indices[1]) == 0:
bounding_boxes.append((0, 0, img.shape[1], img.shape[0]))
cropped_images.append(img)
cropped_masks.append(mask)
combined_cropped_images.append(img)
combined_cropped_masks.append(mask)
else:
min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1]) min_x, max_x = np.min(non_zero_indices[1]), np.max(non_zero_indices[1])
min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0]) min_y, max_y = np.min(non_zero_indices[0]), np.max(non_zero_indices[0])
@ -1934,11 +1948,12 @@ class BatchCropFromMaskAdvanced:
# Resize the cropped image to a fixed size # Resize the cropped image to a fixed size
new_size = max(cropped_img.shape[0], cropped_img.shape[1]) new_size = max(cropped_img.shape[0], cropped_img.shape[1])
resize_transform = Resize(new_size, interpolation = InterpolationMode.NEAREST) resize_transform = Resize(new_size, interpolation=InterpolationMode.NEAREST, max_size=max(img.shape[0], img.shape[1]))
resized_mask = resize_transform(cropped_mask.unsqueeze(0).unsqueeze(0)).squeeze(0).squeeze(0) resized_mask = resize_transform(cropped_mask.unsqueeze(0).unsqueeze(0)).squeeze(0).squeeze(0)
resized_img = resize_transform(cropped_img.permute(2, 0, 1)) resized_img = resize_transform(cropped_img.permute(2, 0, 1))
# Perform the center crop to the desired size # Perform the center crop to the desired size
crop_transform = CenterCrop((self.max_bbox_size, self.max_bbox_size)) # Constrain the crop to the smaller of our bbox or our image so we don't expand past the image dimensions.
crop_transform = CenterCrop((min(self.max_bbox_size, resized_img.shape[1]), min(self.max_bbox_size, resized_img.shape[2])))
cropped_resized_img = crop_transform(resized_img) cropped_resized_img = crop_transform(resized_img)
cropped_images.append(cropped_resized_img.permute(1, 2, 0)) cropped_images.append(cropped_resized_img.permute(1, 2, 0))