feat: Add Mask from ID Node + ColorField Component Improvements

invokeai/app/invocations/image.py

@@ -967,3 +967,56 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
 
         image_dto = context.images.save(image=source_image)
         return ImageOutput.build(image_dto)
+
+
+@invocation(
+    "mask_from_id",
+    title="Mask from ID",
+    tags=["image", "mask", "id"],
+    category="image",
+    version="1.0.0",
+)
+class MaskFromIDInvocation(BaseInvocation, WithMetadata, WithBoard):
+    """Generate a mask for a particular color in an ID Map"""
+
+    image: ImageField = InputField(description="The image to create the mask from")
+    color: ColorField = InputField(description="ID color to mask")
+    threshold: int = InputField(default=100, description="Threshold for color detection")
+    invert: bool = InputField(default=False, description="Whether or not to invert the mask")
+
+    def rgba_to_hex(self, rgba_color: tuple[int, int, int, int]):
+        r, g, b, a = rgba_color
+        hex_code = "#{:02X}{:02X}{:02X}{:02X}".format(r, g, b, int(a * 255))
+        return hex_code
+
+    def id_to_mask(self, id_mask: Image.Image, color: tuple[int, int, int, int], threshold: int = 100):
+        if id_mask.mode != "RGB":
+            id_mask = id_mask.convert("RGB")
+
+        # Can directly just use the tuple but I'll leave this rgba_to_hex here
+        # incase anyone prefers using hex codes directly instead of the color picker
+        hex_color_str = self.rgba_to_hex(color)
+        rgb_color = numpy.array([int(hex_color_str[i : i + 2], 16) for i in (1, 3, 5)])
+
+        # Maybe there's a faster way to calculate this distance but I can't think of any right now.
+        color_distance = numpy.linalg.norm(id_mask - rgb_color, axis=-1)
+
+        # Create a mask based on the threshold and the distance calculated above
+        binary_mask = (color_distance < threshold).astype(numpy.uint8) * 255
+
+        # Convert the mask back to PIL
+        binary_mask_pil = Image.fromarray(binary_mask)
+
+        return binary_mask_pil
+
+    def invoke(self, context: InvocationContext) -> ImageOutput:
+        image = context.images.get_pil(self.image.image_name)
+
+        mask = self.id_to_mask(image, self.color.tuple(), self.threshold)
+
+        if self.invert:
+            mask = ImageOps.invert(mask)
+
+        image_dto = context.images.save(image=mask, image_category=ImageCategory.MASK)
+
+        return ImageOutput.build(image_dto)

invokeai/frontend/web/src/common/util/colorCodeTransformers.ts

@@ -0,0 +1,17 @@
+import type { RgbaColor } from 'react-colorful';
+
+export function rgbaToHex(color: RgbaColor, alpha: boolean = false): string {
+  const hex = ((1 << 24) + (color.r << 16) + (color.g << 8) + color.b).toString(16).slice(1);
+  const alphaHex = Math.round(color.a * 255)
+    .toString(16)
+    .padStart(2, '0');
+  return alpha ? `#${hex}${alphaHex}` : `#${hex}`;
+}
+
+export function hexToRGBA(hex: string, alpha: number) {
+  hex = hex.replace(/^#/, '');
+  const r = parseInt(hex.substring(0, 2), 16);
+  const g = parseInt(hex.substring(2, 4), 16);
+  const b = parseInt(hex.substring(4, 6), 16);
+  return { r, g, b, a: alpha };
+}

invokeai/frontend/web/src/features/nodes/components/flow/nodes/Invocation/fields/inputs/ColorFieldInputComponent.tsx

@@ -1,9 +1,12 @@
+import { Box } from '@invoke-ai/ui-library';
 import { useAppDispatch } from 'app/store/storeHooks';
+import { hexToRGBA, rgbaToHex } from 'common/util/colorCodeTransformers';
+import { colorTokenToCssVar } from 'common/util/colorTokenToCssVar';
 import { fieldColorValueChanged } from 'features/nodes/store/nodesSlice';
 import type { ColorFieldInputInstance, ColorFieldInputTemplate } from 'features/nodes/types/field';
 import { memo, useCallback, useMemo } from 'react';
 import type { RgbaColor } from 'react-colorful';
-import { RgbaColorPicker } from 'react-colorful';
+import { HexColorInput, RgbaColorPicker } from 'react-colorful';
 
 import type { FieldComponentProps } from './types';
 
@@ -26,8 +29,12 @@ const ColorFieldInputComponent = (props: FieldComponentProps<ColorFieldInputInst
   }, [field.value]);
 
   const handleValueChanged = useCallback(
-    (value: RgbaColor) => {
+    (value: RgbaColor | string) => {
       // We need to multiply by 255 to convert from 0-1 to 0-255, which is what the backend needs
+      if (typeof value === 'string') {
+        value = hexToRGBA(value, 1);
+      }
+
       const { r, g, b, a: _a } = value;
       const a = Math.round(_a * 255);
       dispatch(
@@ -41,7 +48,27 @@ const ColorFieldInputComponent = (props: FieldComponentProps<ColorFieldInputInst
     [dispatch, field.name, nodeId]
   );
 
-  return <RgbaColorPicker className="nodrag" color={color} onChange={handleValueChanged} />;
+  return (
+    <Box sx={{ display: 'flex', flexDirection: 'column', gap: 2 }}>
+      <HexColorInput
+        style={{
+          background: colorTokenToCssVar('base.700'),
+          color: colorTokenToCssVar('base.100'),
+          fontSize: 12,
+          paddingInlineStart: 10,
+          borderRadius: 4,
+          paddingBlock: 4,
+          outline: 'none',
+        }}
+        className="nodrag"
+        color={rgbaToHex(color, true)}
+        onChange={handleValueChanged}
+        prefixed
+        alpha
+      />
+      <RgbaColorPicker className="nodrag" color={color} onChange={handleValueChanged} style={{ width: '100%' }} />
+    </Box>
+  );
 };
 
 export default memo(ColorFieldInputComponent);