Skip to content

Commit

Permalink
fix ruff?
Browse files Browse the repository at this point in the history
  • Loading branch information
hipsterusername committed Sep 18, 2024
1 parent 1ea7b92 commit 4367ae0
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 13 deletions.
11 changes: 4 additions & 7 deletions invokeai/app/invocations/image.py
Original file line number Diff line number Diff line change
Expand Up @@ -1074,8 +1074,6 @@ def invoke(self, context: InvocationContext) -> CanvasV2MaskAndCropOutput:
)




@invocation_output("crop_to_object_output")
class CropToObjectOutput(ImageOutput):
offset_top: int = OutputField(description="The number of pixels cropped from the top")
Expand All @@ -1096,9 +1094,8 @@ class CropToObjectInvocation(BaseInvocation, WithMetadata, WithBoard):

image: ImageField = InputField(description="An input mask image with black and white content")
margin: int = InputField(default=0, ge=0, description="The desired margin around the object, as measured in pixels")
object_color: Literal['white', 'black'] = InputField(
default='white',
description="The color of the object to crop around (either 'white' or 'black')"
object_color: Literal["white", "black"] = InputField(
default="white", description="The color of the object to crop around (either 'white' or 'black')"
)

def invoke(self, context: InvocationContext) -> CropToObjectOutput:
Expand All @@ -1112,7 +1109,7 @@ def invoke(self, context: InvocationContext) -> CropToObjectOutput:
np_image = numpy.array(grayscale_image)

# Depending on the object color, find the object pixels
if self.object_color == 'white':
if self.object_color == "white":
# Find white pixels (value > 0)
object_pixels = numpy.argwhere(np_image > 0)
else:
Expand Down Expand Up @@ -1162,4 +1159,4 @@ def invoke(self, context: InvocationContext) -> CropToObjectOutput:
offset_left=offset_left,
offset_right=offset_right,
offset_bottom=offset_bottom,
)
)
9 changes: 3 additions & 6 deletions invokeai/app/invocations/mask.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,15 +108,14 @@ class ImageMaskToTensorInvocation(BaseInvocation, WithMetadata):
def invoke(self, context: InvocationContext) -> MaskOutput:
image = context.images.get_pil(self.image.image_name)
np_image = np.array(image)

# Handle different image modes
if image.mode == 'RGBA':
if image.mode == "RGBA":
alpha_channel = np_image[:, :, 3] # Extract alpha channel
elif image.mode == 'RGB':
elif image.mode == "RGB":
# For RGB images, treat all non-black pixels as opaque.
non_black_mask = np.any(np_image > 0, axis=2) # True for any non-black pixels
alpha_channel = non_black_mask.astype(np.uint8) * 255 # Convert to a mask of 0 or 255
elif image.mode == 'L': # Grayscale images
elif image.mode == "L": # Grayscale images
alpha_channel = np_image # Grayscale image, so we directly use it
else:
raise ValueError(f"Unsupported image mode: {image.mode}")
Expand All @@ -135,8 +134,6 @@ def invoke(self, context: InvocationContext) -> MaskOutput:
)




@invocation(
"tensor_mask_to_image",
title="Tensor Mask to Image",
Expand Down

0 comments on commit 4367ae0

Please sign in to comment.