Skip to content

Commit

Permalink
fix: ✨ mask crop output
Browse files Browse the repository at this point in the history
Follow up of #124
Thanks to @Yurchikian
  • Loading branch information
melMass committed Nov 14, 2023
1 parent e4da832 commit 59a361a
Showing 1 changed file with 8 additions and 4 deletions.
12 changes: 8 additions & 4 deletions nodes/crop.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,19 +141,23 @@ def do_crop(
self, image: torch.Tensor, mask=None, x=0, y=0, width=256, height=256, bbox=None
):
image = image.numpy()
if mask != None:
if mask is not None:
mask = mask.numpy()

if bbox != None:
if bbox is not None:
x, y, width, height = bbox

cropped_image = image[:, y : y + height, x : x + width, :]
cropped_mask = mask[y : y + height, x : x + width] if mask != None else None
cropped_mask = None
if mask is not None:
cropped_mask = (
mask[:, y : y + height, x : x + width] if mask is not None else None
)
crop_data = (x, y, width, height)

return (
torch.from_numpy(cropped_image),
torch.from_numpy(cropped_mask) if mask != None else None,
torch.from_numpy(cropped_mask) if cropped_mask is not None else None,
crop_data,
)

Expand Down

0 comments on commit 59a361a

Please sign in to comment.