From 8911017bd1d8cf00042d081979e7daa36576f6e0 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Wed, 10 Apr 2024 18:56:01 +1000 Subject: [PATCH] feat(ui): selectable & draggable layers --- invokeai/app/invocations/mask.py | 1 + .../src/features/canvas/util/blobToDataURL.ts | 19 +++ .../components/BrushPreview.tsx | 33 ++--- .../components/LayerBoundingBox.tsx | 46 +++++++ .../components/LayerComponent.tsx | 115 ++++++++++++++++++ .../components/LineComponent.tsx | 10 +- .../components/RectComponent.tsx | 14 +-- .../components/RegionalPromptsEditor.tsx | 4 +- .../components/RegionalPromptsStage.tsx | 52 ++++---- .../components/ToolChooser.tsx | 11 +- .../regionalPrompts/hooks/mouseEventHooks.ts | 2 +- .../regionalPrompts/hooks/useTransform.ts | 28 ++--- .../store/regionalPromptsSlice.ts | 77 ++++++++---- .../src/features/regionalPrompts/util/bbox.ts | 97 +++++++++++++++ 14 files changed, 401 insertions(+), 108 deletions(-) create mode 100644 invokeai/frontend/web/src/features/regionalPrompts/components/LayerBoundingBox.tsx create mode 100644 invokeai/frontend/web/src/features/regionalPrompts/components/LayerComponent.tsx create mode 100644 invokeai/frontend/web/src/features/regionalPrompts/util/bbox.ts diff --git a/invokeai/app/invocations/mask.py b/invokeai/app/invocations/mask.py index 72c5886336..fcc5b04caf 100644 --- a/invokeai/app/invocations/mask.py +++ b/invokeai/app/invocations/mask.py @@ -48,6 +48,7 @@ class AlphaMaskToTensorInvocation(BaseInvocation): """Convert a mask image to a tensor. Opaque regions are 1 and transparent regions are 0.""" image: ImageField = InputField(description="The mask image to convert.") + invert: bool = InputField(default=False, description="Invert the mask (1s become 0s and 0s become 1s).") def invoke(self, context: InvocationContext) -> MaskOutput: image = context.images.get_pil(self.image.image_name) diff --git a/invokeai/frontend/web/src/features/canvas/util/blobToDataURL.ts b/invokeai/frontend/web/src/features/canvas/util/blobToDataURL.ts index 2443396105..f29010c99c 100644 --- a/invokeai/frontend/web/src/features/canvas/util/blobToDataURL.ts +++ b/invokeai/frontend/web/src/features/canvas/util/blobToDataURL.ts @@ -7,3 +7,22 @@ export const blobToDataURL = (blob: Blob): Promise => { reader.readAsDataURL(blob); }); }; + +export function imageDataToDataURL(imageData: ImageData): string { + const { width, height } = imageData; + + // Create a canvas to transfer the ImageData to + const canvas = document.createElement('canvas'); + canvas.width = width; + canvas.height = height; + + // Draw the ImageData onto the canvas + const ctx = canvas.getContext('2d'); + if (!ctx) { + throw new Error('Unable to get canvas context'); + } + ctx.putImageData(imageData, 0, 0); + + // Convert the canvas to a data URL (base64) + return canvas.toDataURL(); +} diff --git a/invokeai/frontend/web/src/features/regionalPrompts/components/BrushPreview.tsx b/invokeai/frontend/web/src/features/regionalPrompts/components/BrushPreview.tsx index 759be3f1ce..2792d4ae81 100644 --- a/invokeai/frontend/web/src/features/regionalPrompts/components/BrushPreview.tsx +++ b/invokeai/frontend/web/src/features/regionalPrompts/components/BrushPreview.tsx @@ -4,8 +4,9 @@ import { rgbColorToString } from 'features/canvas/util/colorToString'; import { $cursorPosition } from 'features/regionalPrompts/store/regionalPromptsSlice'; import { Circle, Group } from 'react-konva'; -export const BrushPreviewFill = () => { +const useBrushData = () => { const brushSize = useAppSelector((s) => s.regionalPrompts.brushSize); + const tool = useAppSelector((s) => s.regionalPrompts.tool); const color = useAppSelector((s) => { const _color = s.regionalPrompts.layers.find((l) => l.id === s.regionalPrompts.selectedLayer)?.color; if (!_color) { @@ -15,25 +16,29 @@ export const BrushPreviewFill = () => { }); const pos = useStore($cursorPosition); - if (!brushSize || !color || !pos) { + return { brushSize, tool, color, pos }; +}; + +export const BrushPreviewFill = () => { + const { brushSize, tool, color, pos } = useBrushData(); + if (!brushSize || !color || !pos || tool === 'move') { return null; } - return ; + return ( + + ); }; export const BrushPreviewOutline = () => { - const brushSize = useAppSelector((s) => s.regionalPrompts.brushSize); - const color = useAppSelector((s) => { - const _color = s.regionalPrompts.layers.find((l) => l.id === s.regionalPrompts.selectedLayer)?.color; - if (!_color) { - return null; - } - return rgbColorToString(_color); - }); - const pos = useStore($cursorPosition); - - if (!brushSize || !color || !pos) { + const { brushSize, tool, color, pos } = useBrushData(); + if (!brushSize || !color || !pos || tool === 'move') { return null; } diff --git a/invokeai/frontend/web/src/features/regionalPrompts/components/LayerBoundingBox.tsx b/invokeai/frontend/web/src/features/regionalPrompts/components/LayerBoundingBox.tsx new file mode 100644 index 0000000000..2745f4eea5 --- /dev/null +++ b/invokeai/frontend/web/src/features/regionalPrompts/components/LayerBoundingBox.tsx @@ -0,0 +1,46 @@ +import { createSelector } from '@reduxjs/toolkit'; +import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; +import { layerSelected, selectRegionalPromptsSlice } from 'features/regionalPrompts/store/regionalPromptsSlice'; +import { useCallback, useMemo } from 'react'; +import { Rect as KonvaRect } from 'react-konva'; + +type Props = { + layerId: string; +}; + +export const LayerBoundingBox = ({ layerId }: Props) => { + const dispatch = useAppDispatch(); + const tool = useAppSelector((s) => s.regionalPrompts.tool); + const selectedLayer = useAppSelector((s) => s.regionalPrompts.selectedLayer); + + const onMouseDown = useCallback(() => { + dispatch(layerSelected(layerId)); + }, [dispatch, layerId]); + + const selectBbox = useMemo( + () => + createSelector( + selectRegionalPromptsSlice, + (regionalPrompts) => regionalPrompts.layers.find((layer) => layer.id === layerId)?.bbox ?? null + ), + [layerId] + ); + const bbox = useAppSelector(selectBbox); + + if (!bbox || tool !== 'move') { + return null; + } + + return ( + + ); +}; diff --git a/invokeai/frontend/web/src/features/regionalPrompts/components/LayerComponent.tsx b/invokeai/frontend/web/src/features/regionalPrompts/components/LayerComponent.tsx new file mode 100644 index 0000000000..030df4943f --- /dev/null +++ b/invokeai/frontend/web/src/features/regionalPrompts/components/LayerComponent.tsx @@ -0,0 +1,115 @@ +import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; +import getScaledCursorPosition from 'features/canvas/util/getScaledCursorPosition'; +import { BrushPreviewFill } from 'features/regionalPrompts/components/BrushPreview'; +import { LayerBoundingBox } from 'features/regionalPrompts/components/LayerBoundingBox'; +import { LineComponent } from 'features/regionalPrompts/components/LineComponent'; +import { RectComponent } from 'features/regionalPrompts/components/RectComponent'; +import { useLayer } from 'features/regionalPrompts/hooks/layerStateHooks'; +import { $stage, layerBboxChanged, layerTranslated } from 'features/regionalPrompts/store/regionalPromptsSlice'; +import { getKonvaLayerBbox } from 'features/regionalPrompts/util/bbox'; +import type { Group as KonvaGroupType } from 'konva/lib/Group'; +import type { Layer as KonvaLayerType } from 'konva/lib/Layer'; +import type { KonvaEventObject, Node as KonvaNodeType, NodeConfig as KonvaNodeConfigType } from 'konva/lib/Node'; +import type { IRect, Vector2d } from 'konva/lib/types'; +import type React from 'react'; +import { useCallback, useEffect, useRef } from 'react'; +import { Group as KonvaGroup, Layer as KonvaLayer } from 'react-konva'; + +type Props = { + id: string; +}; + +const filterChildren = (item: KonvaNodeType) => item.name() !== 'regionalPromptLayerObjectGroup'; + +export const LayerComponent: React.FC = ({ id }) => { + const dispatch = useAppDispatch(); + const layer = useLayer(id); + const selectedLayer = useAppSelector((s) => s.regionalPrompts.selectedLayer); + const tool = useAppSelector((s) => s.regionalPrompts.tool); + const layerRef = useRef(null); + const groupRef = useRef(null); + + const onChangeBbox = useCallback( + (bbox: IRect | null) => { + dispatch(layerBboxChanged({ layerId: layer.id, bbox })); + }, + [dispatch, layer.id] + ); + + const onDragEnd = useCallback( + (e: KonvaEventObject) => { + dispatch(layerTranslated({ layerId: id, x: e.target.x(), y: e.target.y() })); + }, + [dispatch, id] + ); + + const onDragMove = useCallback( + (e: KonvaEventObject) => { + dispatch(layerTranslated({ layerId: id, x: e.target.x(), y: e.target.y() })); + }, + [dispatch, id] + ); + + const dragBoundFunc = useCallback(function (this: KonvaNodeType, pos: Vector2d) { + const stage = $stage.get(); + if (!stage) { + return this.getAbsolutePosition(); + } + const cursorPos = getScaledCursorPosition(stage); + if (!cursorPos) { + return this.getAbsolutePosition(); + } + if (cursorPos.x < 0 || cursorPos.x > stage.width() || cursorPos.y < 0 || cursorPos.y > stage.height()) { + return this.getAbsolutePosition(); + } + + return pos; + }, []); + + useEffect(() => { + if (!layerRef.current || tool !== 'move') { + return; + } + if (layer.objects.length === 0) { + onChangeBbox(null); + return; + } + onChangeBbox(getKonvaLayerBbox(layerRef.current, filterChildren)); + }, [tool, layer.objects, onChangeBbox]); + + if (!layer.isVisible) { + return null; + } + + return ( + <> + + + {layer.objects.map((obj) => { + if (obj.kind === 'line') { + return ; + } + if (obj.kind === 'fillRect') { + return ; + } + })} + + + + {layer.id === selectedLayer && } + + ); +}; diff --git a/invokeai/frontend/web/src/features/regionalPrompts/components/LineComponent.tsx b/invokeai/frontend/web/src/features/regionalPrompts/components/LineComponent.tsx index eaecc7625d..a0250ab6da 100644 --- a/invokeai/frontend/web/src/features/regionalPrompts/components/LineComponent.tsx +++ b/invokeai/frontend/web/src/features/regionalPrompts/components/LineComponent.tsx @@ -1,20 +1,18 @@ import { rgbColorToString } from 'features/canvas/util/colorToString'; -import { useTransform } from 'features/regionalPrompts/hooks/useTransform'; import type { LineObject } from 'features/regionalPrompts/store/regionalPromptsSlice'; import type { RgbColor } from 'react-colorful'; import { Line } from 'react-konva'; type Props = { + layerId: string; line: LineObject; color: RgbColor; }; -export const LineComponent = ({ line, color }: Props) => { - const { shapeRef } = useTransform(line); - +export const LineComponent = ({ layerId, line, color }: Props) => { return ( { lineCap="round" lineJoin="round" shadowForStrokeEnabled={false} - listening={false} globalCompositeOperation={line.tool === 'brush' ? 'source-over' : 'destination-out'} + listening={false} /> ); }; diff --git a/invokeai/frontend/web/src/features/regionalPrompts/components/RectComponent.tsx b/invokeai/frontend/web/src/features/regionalPrompts/components/RectComponent.tsx index 1401731dd2..a8248e1301 100644 --- a/invokeai/frontend/web/src/features/regionalPrompts/components/RectComponent.tsx +++ b/invokeai/frontend/web/src/features/regionalPrompts/components/RectComponent.tsx @@ -1,5 +1,4 @@ import { rgbColorToString } from 'features/canvas/util/colorToString'; -import { useTransform } from 'features/regionalPrompts/hooks/useTransform'; import type { FillRectObject } from 'features/regionalPrompts/store/regionalPromptsSlice'; import type { RgbColor } from 'react-colorful'; import { Rect } from 'react-konva'; @@ -10,18 +9,7 @@ type Props = { }; export const RectComponent = ({ rect, color }: Props) => { - const { shapeRef } = useTransform(rect); - return ( - + ); }; diff --git a/invokeai/frontend/web/src/features/regionalPrompts/components/RegionalPromptsEditor.tsx b/invokeai/frontend/web/src/features/regionalPrompts/components/RegionalPromptsEditor.tsx index 23491b9efa..b39ab6587b 100644 --- a/invokeai/frontend/web/src/features/regionalPrompts/components/RegionalPromptsEditor.tsx +++ b/invokeai/frontend/web/src/features/regionalPrompts/components/RegionalPromptsEditor.tsx @@ -9,6 +9,7 @@ import { RegionalPromptsStage } from 'features/regionalPrompts/components/Region import { ToolChooser } from 'features/regionalPrompts/components/ToolChooser'; import { selectRegionalPromptsSlice } from 'features/regionalPrompts/store/regionalPromptsSlice'; import { getLayerBlobs } from 'features/regionalPrompts/util/getLayerBlobs'; +import { ImageSizeLinear } from 'features/settingsAccordions/components/ImageSettingsAccordion/ImageSizeLinear'; const selectLayerIdsReversed = createSelector(selectRegionalPromptsSlice, (regionalPrompts) => regionalPrompts.layers.map((l) => l.id).reverse() @@ -18,10 +19,11 @@ export const RegionalPromptsEditor = () => { const layerIdsReversed = useAppSelector(selectLayerIdsReversed); return ( - + + {layerIdsReversed.map((id) => ( diff --git a/invokeai/frontend/web/src/features/regionalPrompts/components/RegionalPromptsStage.tsx b/invokeai/frontend/web/src/features/regionalPrompts/components/RegionalPromptsStage.tsx index 5a25bf8d1b..b231db65a1 100644 --- a/invokeai/frontend/web/src/features/regionalPrompts/components/RegionalPromptsStage.tsx +++ b/invokeai/frontend/web/src/features/regionalPrompts/components/RegionalPromptsStage.tsx @@ -1,9 +1,8 @@ import { chakra } from '@invoke-ai/ui-library'; import { createSelector } from '@reduxjs/toolkit'; import { useAppSelector } from 'app/store/storeHooks'; -import { BrushPreviewFill, BrushPreviewOutline } from 'features/regionalPrompts/components/BrushPreview'; -import { LineComponent } from 'features/regionalPrompts/components/LineComponent'; -import { RectComponent } from 'features/regionalPrompts/components/RectComponent'; +import { BrushPreviewOutline } from 'features/regionalPrompts/components/BrushPreview'; +import { LayerComponent } from 'features/regionalPrompts/components/LayerComponent'; import { useMouseDown, useMouseEnter, @@ -13,25 +12,23 @@ import { } from 'features/regionalPrompts/hooks/mouseEventHooks'; import { $stage, selectRegionalPromptsSlice } from 'features/regionalPrompts/store/regionalPromptsSlice'; import type Konva from 'konva'; -import { memo, useCallback, useRef } from 'react'; +import { memo, useCallback, useMemo, useRef } from 'react'; import { Layer, Stage } from 'react-konva'; -const selectVisibleLayers = createSelector(selectRegionalPromptsSlice, (regionalPrompts) => - regionalPrompts.layers.filter((l) => l.isVisible) +const selectLayerIds = createSelector(selectRegionalPromptsSlice, (regionalPrompts) => + regionalPrompts.layers.map((l) => l.id) ); const ChakraStage = chakra(Stage, { shouldForwardProp: (prop) => !['sx'].includes(prop), }); -const stageSx = { - border: '1px solid green', -}; - export const RegionalPromptsStage: React.FC = memo(() => { - const layers = useAppSelector(selectVisibleLayers); - const selectedLayer = useAppSelector((s) => s.regionalPrompts.selectedLayer); + const layerIds = useAppSelector(selectLayerIds); const stageRef = useRef(null); + const width = useAppSelector((s) => s.generation.width); + const height = useAppSelector((s) => s.generation.height); + const tool = useAppSelector((s) => s.regionalPrompts.tool); const onMouseDown = useMouseDown(stageRef); const onMouseUp = useMouseUp(stageRef); const onMouseMove = useMouseMove(stageRef); @@ -41,34 +38,33 @@ export const RegionalPromptsStage: React.FC = memo(() => { $stage.set(el); stageRef.current = el; }, []); + const sx = useMemo( + () => ({ + border: '1px solid cyan', + cursor: tool === 'move' ? 'default' : 'none', + }), + [tool] + ); return ( - {layers.map((layer) => ( - - {layer.objects.map((obj) => { - if (obj.kind === 'line') { - return ; - } - if (obj.kind === 'fillRect') { - return ; - } - })} - {layer.id === selectedLayer && } - + {layerIds.map((id) => ( + ))} - + diff --git a/invokeai/frontend/web/src/features/regionalPrompts/components/ToolChooser.tsx b/invokeai/frontend/web/src/features/regionalPrompts/components/ToolChooser.tsx index 4d178548ec..daf039defb 100644 --- a/invokeai/frontend/web/src/features/regionalPrompts/components/ToolChooser.tsx +++ b/invokeai/frontend/web/src/features/regionalPrompts/components/ToolChooser.tsx @@ -2,7 +2,7 @@ import { ButtonGroup, IconButton } from '@invoke-ai/ui-library'; import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; import { toolChanged } from 'features/regionalPrompts/store/regionalPromptsSlice'; import { useCallback } from 'react'; -import { PiEraserBold, PiPaintBrushBold } from 'react-icons/pi'; +import { PiArrowsOutCardinalBold, PiEraserBold, PiPaintBrushBold } from 'react-icons/pi'; export const ToolChooser: React.FC = () => { const tool = useAppSelector((s) => s.regionalPrompts.tool); @@ -13,6 +13,9 @@ export const ToolChooser: React.FC = () => { const setToolToEraser = useCallback(() => { dispatch(toolChanged('eraser')); }, [dispatch]); + const setToolToMove = useCallback(() => { + dispatch(toolChanged('move')); + }, [dispatch]); return ( @@ -28,6 +31,12 @@ export const ToolChooser: React.FC = () => { variant={tool === 'eraser' ? 'solid' : 'outline'} onClick={setToolToEraser} /> + } + variant={tool === 'move' ? 'solid' : 'outline'} + onClick={setToolToMove} + /> ); }; diff --git a/invokeai/frontend/web/src/features/regionalPrompts/hooks/mouseEventHooks.ts b/invokeai/frontend/web/src/features/regionalPrompts/hooks/mouseEventHooks.ts index 084bb24294..17e5cbff7b 100644 --- a/invokeai/frontend/web/src/features/regionalPrompts/hooks/mouseEventHooks.ts +++ b/invokeai/frontend/web/src/features/regionalPrompts/hooks/mouseEventHooks.ts @@ -13,7 +13,7 @@ import type { KonvaEventObject } from 'konva/lib/Node'; import type { MutableRefObject } from 'react'; import { useCallback } from 'react'; -const getTool = () => getStore().getState().regionalPrompts.tool; +export const getTool = () => getStore().getState().regionalPrompts.tool; const getIsFocused = (stage: Konva.Stage) => { return stage.container().contains(document.activeElement); diff --git a/invokeai/frontend/web/src/features/regionalPrompts/hooks/useTransform.ts b/invokeai/frontend/web/src/features/regionalPrompts/hooks/useTransform.ts index d04cbb8c55..87d4e3b9d4 100644 --- a/invokeai/frontend/web/src/features/regionalPrompts/hooks/useTransform.ts +++ b/invokeai/frontend/web/src/features/regionalPrompts/hooks/useTransform.ts @@ -1,30 +1,18 @@ -import type { FillRectObject, LayerObject, LineObject } from 'features/regionalPrompts/store/regionalPromptsSlice'; -import type { Image } from 'konva/lib/shapes/Image'; -import type { Line } from 'konva/lib/shapes/Line'; -import type { Rect } from 'konva/lib/shapes/Rect'; -import type { Transformer } from 'konva/lib/shapes/Transformer'; +import type { Group as KonvaGroupType } from 'konva/lib/Group'; +import type { Transformer as KonvaTransformerType } from 'konva/lib/shapes/Transformer'; import { useEffect, useRef } from 'react'; -type ShapeType = T extends LineObject ? Line : T extends FillRectObject ? Rect : Image; - -export const useTransform = (object: TObject) => { - const shapeRef = useRef>(null); - const transformerRef = useRef(null); +export const useTransform = () => { + const shapeRef = useRef(null); + const transformerRef = useRef(null); useEffect(() => { - if (!object.isSelected) { - return; - } - if (!transformerRef.current || !shapeRef.current) { return; } - - if (object.isSelected) { - transformerRef.current.nodes([shapeRef.current]); - transformerRef.current.getLayer()?.batchDraw(); - } - }, [object.isSelected]); + transformerRef.current.nodes([shapeRef.current]); + transformerRef.current.getLayer()?.batchDraw(); + }, []); return { shapeRef, transformerRef }; }; diff --git a/invokeai/frontend/web/src/features/regionalPrompts/store/regionalPromptsSlice.ts b/invokeai/frontend/web/src/features/regionalPrompts/store/regionalPromptsSlice.ts index 0c644f8808..faafffa936 100644 --- a/invokeai/frontend/web/src/features/regionalPrompts/store/regionalPromptsSlice.ts +++ b/invokeai/frontend/web/src/features/regionalPrompts/store/regionalPromptsSlice.ts @@ -3,17 +3,16 @@ import { createSlice } from '@reduxjs/toolkit'; import type { PersistConfig, RootState } from 'app/store/store'; import { moveBackward, moveForward, moveToBack, moveToFront } from 'common/util/arrayUtils'; import type Konva from 'konva'; -import type { Vector2d } from 'konva/lib/types'; +import type { IRect, Vector2d } from 'konva/lib/types'; import { atom } from 'nanostores'; import type { RgbColor } from 'react-colorful'; import { assert } from 'tsafe'; import { v4 as uuidv4 } from 'uuid'; -export type Tool = 'brush' | 'eraser'; +export type Tool = 'brush' | 'eraser' | 'move'; type LayerObjectBase = { id: string; - isSelected: boolean; }; type ImageObject = LayerObjectBase & { @@ -45,6 +44,9 @@ export type LayerObject = ImageObject | LineObject | FillRectObject; type LayerBase = { id: string; isVisible: boolean; + x: number; + y: number; + bbox: IRect | null; }; type PromptRegionLayer = LayerBase & { @@ -54,7 +56,7 @@ type PromptRegionLayer = LayerBase & { color: RgbColor; }; -type Layer = PromptRegionLayer; +export type Layer = PromptRegionLayer; type RegionalPromptsState = { _version: 1; @@ -81,7 +83,7 @@ export const regionalPromptsSlice = createSlice({ layerAdded: { reducer: (state, action: PayloadAction) => { const newLayer = buildLayer(action.meta.id, action.payload, state.layers.length); - state.layers.unshift(newLayer); + state.layers.push(newLayer); state.selectedLayer = newLayer.id; }, prepare: (payload: Layer['kind']) => ({ payload, meta: { id: uuidv4() } }), @@ -102,6 +104,9 @@ export const regionalPromptsSlice = createSlice({ return; } layer.objects = []; + layer.bbox = null; + layer.isVisible = true; + layer.prompt = ''; }, layerDeleted: (state, action: PayloadAction) => { state.layers = state.layers.filter((l) => l.id !== action.payload); @@ -125,6 +130,23 @@ export const regionalPromptsSlice = createSlice({ // Because the layers are in reverse order, moving to the back is equivalent to moving to the front moveToFront(state.layers, cb); }, + layerTranslated: (state, action: PayloadAction<{ layerId: string; x: number; y: number }>) => { + const { layerId, x, y } = action.payload; + const layer = state.layers.find((l) => l.id === layerId); + if (!layer) { + return; + } + layer.x = x; + layer.y = y; + }, + layerBboxChanged: (state, action: PayloadAction<{ layerId: string; bbox: IRect | null }>) => { + const { layerId, bbox } = action.payload; + const layer = state.layers.find((l) => l.id === layerId); + if (!layer) { + return; + } + layer.bbox = bbox; + }, promptChanged: (state, action: PayloadAction<{ layerId: string; prompt: string }>) => { const { layerId, prompt } = action.payload; const layer = state.layers.find((l) => l.id === layerId); @@ -142,25 +164,31 @@ export const regionalPromptsSlice = createSlice({ layer.color = color; }, lineAdded: { - reducer: (state, action: PayloadAction) => { - const selectedLayer = state.layers.find((l) => l.id === state.selectedLayer); - if (!selectedLayer || selectedLayer.kind !== 'promptRegionLayer') { + reducer: (state, action: PayloadAction<[number, number], string, { id: string }>) => { + const layer = state.layers.find((l) => l.id === state.selectedLayer); + if (!layer || layer.kind !== 'promptRegionLayer') { return; } - selectedLayer.objects.push(buildLine(action.meta.id, action.payload, state.brushSize, state.tool)); + layer.objects.push({ + kind: 'line', + tool: state.tool, + id: action.meta.id, + points: [action.payload[0] - layer.x, action.payload[1] - layer.y], + strokeWidth: state.brushSize, + }); }, - prepare: (payload: number[]) => ({ payload, meta: { id: uuidv4() } }), + prepare: (payload: [number, number]) => ({ payload, meta: { id: uuidv4() } }), }, - pointsAdded: (state, action: PayloadAction) => { - const selectedLayer = state.layers.find((l) => l.id === state.selectedLayer); - if (!selectedLayer || selectedLayer.kind !== 'promptRegionLayer') { + pointsAdded: (state, action: PayloadAction<[number, number]>) => { + const layer = state.layers.find((l) => l.id === state.selectedLayer); + if (!layer || layer.kind !== 'promptRegionLayer') { return; } - const lastLine = selectedLayer.objects.findLast(isLine); + const lastLine = layer.objects.findLast(isLine); if (!lastLine) { return; } - lastLine.points.push(...action.payload); + lastLine.points.push(action.payload[0] - layer.x, action.payload[1] - layer.y); }, brushSizeChanged: (state, action: PayloadAction) => { state.brushSize = action.payload; @@ -187,24 +215,18 @@ const buildLayer = (id: string, kind: Layer['kind'], layerCount: number): Layer return { id, isVisible: true, + bbox: null, kind, prompt: '', objects: [], color, + x: 0, + y: 0, }; } assert(false, `Unknown layer kind: ${kind}`); }; -const buildLine = (id: string, points: number[], brushSize: number, tool: Tool): LineObject => ({ - isSelected: false, - kind: 'line', - tool, - id, - points, - strokeWidth: brushSize, -}); - export const { layerAdded, layerSelected, @@ -221,6 +243,8 @@ export const { layerMovedBackward, layerMovedToBack, toolChanged, + layerTranslated, + layerBboxChanged, } = regionalPromptsSlice.actions; export const selectRegionalPromptsSlice = (state: RootState) => state.regionalPrompts; @@ -241,3 +265,8 @@ export const $isMouseDown = atom(false); export const $isMouseOver = atom(false); export const $cursorPosition = atom(null); export const $stage = atom(null); +export const getStage = (): Konva.Stage => { + const stage = $stage.get(); + assert(stage); + return stage; +}; diff --git a/invokeai/frontend/web/src/features/regionalPrompts/util/bbox.ts b/invokeai/frontend/web/src/features/regionalPrompts/util/bbox.ts new file mode 100644 index 0000000000..eceac762a3 --- /dev/null +++ b/invokeai/frontend/web/src/features/regionalPrompts/util/bbox.ts @@ -0,0 +1,97 @@ +import Konva from 'konva'; +import type { Layer as KonvaLayerType } from 'konva/lib/Layer'; +import type { Node as KonvaNodeType, NodeConfig as KonvaNodeConfigType } from 'konva/lib/Node'; +import type { IRect } from 'konva/lib/types'; +import { assert } from 'tsafe'; + +/** + * Get the bounding box of an image. + * @param imageData The ImageData object to get the bounding box of. + * @returns The minimum and maximum x and y values of the image's bounding box. + */ +export const getImageDataBbox = (imageData: ImageData) => { + const { data, width, height } = imageData; + let minX = width; + let minY = height; + let maxX = 0; + let maxY = 0; + + for (let y = 0; y < height; y++) { + for (let x = 0; x < width; x++) { + const alpha = data[(y * width + x) * 4 + 3] ?? 0; + if (alpha > 0) { + minX = Math.min(minX, x); + maxX = Math.max(maxX, x); + minY = Math.min(minY, y); + maxY = Math.max(maxY, y); + } + } + } + + return { minX, minY, maxX, maxY }; +}; + +/** + * Get the bounding box of a regional prompt konva layer. This function has special handling for regional prompt layers. + * @param layer The konva layer to get the bounding box of. + * @param filterChildren Optional filter function to exclude certain children from the bounding box calculation. Defaults to including all children. + */ +export const getKonvaLayerBbox = ( + layer: KonvaLayerType, + filterChildren?: (item: KonvaNodeType) => boolean +): IRect => { + // To calculate the layer's bounding box, we must first render it to a pixel array, then do some math. + // We can't use konva's `layer.getClientRect()`, because this includes all shapes, not just visible ones. + // That would include eraser strokes, and the resultant bbox would be too large. + const stage = layer.getStage(); + + // Construct and offscreen canvas and add just the layer to it. + const offscreenStageContainer = document.createElement('div'); + const offscreenStage = new Konva.Stage({ + container: offscreenStageContainer, + width: stage.width(), + height: stage.height(), + }); + + // Clone the layer and filter out unwanted children. + // TODO: Would be more efficient to create a totally new layer and add only the children we want, but possibly less + // accurate, as we wouldn't get the original layer's config and such. + const layerClone = layer.clone(); + if (filterChildren) { + for (const child of layerClone.getChildren(filterChildren)) { + child.destroy(); + } + } + + offscreenStage.add(layerClone.clone()); + + // Get the layer's image data, ensuring we capture an area large enough to include the full layer, including any + // portions that are outside the current stage bounds. + const layerRect = layerClone.getClientRect(); + + // Render the canvas, large enough to capture the full layer. + const x = -layerRect.width; // start from left of layer, as far left as the layer might be + const y = -layerRect.height; // start from top of layer, as far up as the layer might be + const width = stage.width() + layerRect.width * 2; // stage width + layer width on left/right + const height = stage.height() + layerRect.height * 2; // stage height + layer height on top/bottom + + // Capture the image data with the above rect. + const layerImageData = offscreenStage + .toCanvas({ x, y, width, height }) + .getContext('2d') + ?.getImageData(0, 0, width, height); + assert(layerImageData, "Unable to get layer's image data"); + + // Calculate the layer's bounding box. + const layerBbox = getImageDataBbox(layerImageData); + + // Correct the bounding box to be relative to the layer's position. + const correctedLayerBbox = { + x: layerBbox.minX - layerRect.width - layer.x(), + y: layerBbox.minY - layerRect.height - layer.y(), + width: layerBbox.maxX - layerBbox.minX, + height: layerBbox.maxY - layerBbox.minY, + }; + + return correctedLayerBbox; +};