Skip to content

Commit 259304b

Browse files
DustyShoelstein
andauthored
Feature(UI): add extract masked area from raster layers (#8667)
* chore: localize extraction errors * chore: rename extract masked area menu item * chore: rename inpaint mask extract component * fix: use mask bounds for extraction region * Prettier format applied to InpaintMaskMenuItemsExtractMaskedArea.tsx * Fix base64 image import bug in extracted area in InpaintMaskMenuItemsExtractMaskedArea.tsx and removed unused locales entries in en.json * Fix formatting issue in InpaintMaskMenuItemsExtractMaskedArea.tsx * Minor comment fix --------- Co-authored-by: Lincoln Stein <[email protected]>
1 parent 2be701c commit 259304b

File tree

3 files changed

+155
-0
lines changed

3 files changed

+155
-0
lines changed

invokeai/frontend/web/public/locales/en.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2152,6 +2152,7 @@
21522152
"clearCaches": "Clear Caches",
21532153
"recalculateRects": "Recalculate Rects",
21542154
"clipToBbox": "Clip Strokes to Bbox",
2155+
"extractRegion": "Extract Region",
21552156
"outputOnlyMaskedRegions": "Output Only Generated Regions",
21562157
"addLayer": "Add Layer",
21572158
"duplicate": "Duplicate",

invokeai/frontend/web/src/features/controlLayers/components/InpaintMask/InpaintMaskMenuItems.tsx

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import { CanvasEntityMenuItemsTransform } from 'features/controlLayers/component
1010
import { InpaintMaskMenuItemsAddModifiers } from 'features/controlLayers/components/InpaintMask/InpaintMaskMenuItemsAddModifiers';
1111
import { InpaintMaskMenuItemsConvertToSubMenu } from 'features/controlLayers/components/InpaintMask/InpaintMaskMenuItemsConvertToSubMenu';
1212
import { InpaintMaskMenuItemsCopyToSubMenu } from 'features/controlLayers/components/InpaintMask/InpaintMaskMenuItemsCopyToSubMenu';
13+
import { InpaintMaskMenuItemsExtractMaskedArea } from 'features/controlLayers/components/InpaintMask/InpaintMaskMenuItemsExtractMaskedArea';
1314
import { memo } from 'react';
1415

1516
export const InpaintMaskMenuItems = memo(() => {
@@ -24,6 +25,7 @@ export const InpaintMaskMenuItems = memo(() => {
2425
<InpaintMaskMenuItemsAddModifiers />
2526
<MenuDivider />
2627
<CanvasEntityMenuItemsTransform />
28+
<InpaintMaskMenuItemsExtractMaskedArea />
2729
<MenuDivider />
2830
<CanvasEntityMenuItemsMergeDown />
2931
<InpaintMaskMenuItemsCopyToSubMenu />
Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
import { MenuItem } from '@invoke-ai/ui-library';
2+
import { logger } from 'app/logging/logger';
3+
import { useCanvasManager } from 'features/controlLayers/contexts/CanvasManagerProviderGate';
4+
import { useEntityIdentifierContext } from 'features/controlLayers/contexts/EntityIdentifierContext';
5+
import { useCanvasIsBusy } from 'features/controlLayers/hooks/useCanvasIsBusy';
6+
import { canvasToBlob, canvasToImageData } from 'features/controlLayers/konva/util';
7+
import type { CanvasImageState, Rect } from 'features/controlLayers/store/types';
8+
import { imageDTOToImageObject } from 'features/controlLayers/store/util';
9+
import { toast } from 'features/toast/toast';
10+
import { memo, useCallback } from 'react';
11+
import { useTranslation } from 'react-i18next';
12+
import { PiSelectionBackgroundBold } from 'react-icons/pi';
13+
import { serializeError } from 'serialize-error';
14+
import { uploadImage } from 'services/api/endpoints/images';
15+
16+
const log = logger('canvas');
17+
18+
export const InpaintMaskMenuItemsExtractMaskedArea = memo(() => {
19+
const canvasManager = useCanvasManager();
20+
const entityIdentifier = useEntityIdentifierContext('inpaint_mask');
21+
const isBusy = useCanvasIsBusy();
22+
const { t } = useTranslation();
23+
24+
const onExtract = useCallback(() => {
25+
void (async () => {
26+
// The active inpaint mask layer is required to build the mask used for extraction.
27+
const maskAdapter = canvasManager.getAdapter(entityIdentifier);
28+
if (!maskAdapter) {
29+
log.error({ entityIdentifier }, 'Inpaint mask adapter not found when extracting masked area');
30+
toast({ status: 'error', title: 'Unable to extract masked area.' });
31+
return;
32+
}
33+
34+
try {
35+
// Use the mask's bounding box in stage coordinates to constrain the extraction region.
36+
const maskPixelRect = maskAdapter.transformer.$pixelRect.get();
37+
const maskPosition = maskAdapter.state.position;
38+
const rect: Rect = {
39+
x: Math.floor(maskPosition.x + maskPixelRect.x),
40+
y: Math.floor(maskPosition.y + maskPixelRect.y),
41+
width: Math.floor(maskPixelRect.width),
42+
height: Math.floor(maskPixelRect.height),
43+
};
44+
45+
// Abort when the canvas is effectively empty—no pixels to extract.
46+
if (rect.width <= 0 || rect.height <= 0) {
47+
toast({ status: 'warning', title: 'Canvas is empty.' });
48+
return;
49+
}
50+
51+
// Gather the visible raster layer adapters so we can composite them into a single bitmap.
52+
const rasterAdapters = canvasManager.compositor.getVisibleAdaptersOfType('raster_layer');
53+
54+
let compositeImageData: ImageData;
55+
if (rasterAdapters.length === 0) {
56+
// No visible raster layers—create a transparent buffer that matches the canvas bounds.
57+
compositeImageData = new ImageData(rect.width, rect.height);
58+
} else {
59+
// Render the visible raster layers into an offscreen canvas restricted to the canvas bounds.
60+
const compositeCanvas = canvasManager.compositor.getCompositeCanvas(rasterAdapters, rect);
61+
compositeImageData = canvasToImageData(compositeCanvas);
62+
}
63+
64+
// Render the inpaint mask layer into a canvas so we have the alpha data that defines the mask.
65+
const maskCanvas = maskAdapter.getCanvas(rect);
66+
const maskImageData = canvasToImageData(maskCanvas);
67+
68+
if (maskImageData.width !== compositeImageData.width || maskImageData.height !== compositeImageData.height) {
69+
// Bail out if the mask and composite buffers disagree on dimensions.
70+
log.error(
71+
{
72+
maskDimensions: { width: maskImageData.width, height: maskImageData.height },
73+
compositeDimensions: { width: compositeImageData.width, height: compositeImageData.height },
74+
},
75+
'Mask and composite dimensions did not match when extracting masked area'
76+
);
77+
toast({ status: 'error', title: 'Unable to extract masked area.' });
78+
return;
79+
}
80+
81+
const compositeArray = compositeImageData.data;
82+
const maskArray = maskImageData.data;
83+
84+
if (!compositeArray || !maskArray) {
85+
toast({ status: 'error', title: 'Cannot extract: image or mask data is missing.' });
86+
return;
87+
}
88+
89+
const outputArray = new Uint8ClampedArray(compositeArray.length);
90+
91+
// Apply the mask alpha channel to each pixel in the composite, keeping RGB untouched and only masking alpha.
92+
for (let i = 0; i < compositeArray.length; i += 4) {
93+
const maskAlpha = (maskArray[i + 3] ?? 0) / 255 || 0;
94+
outputArray[i] = compositeArray[i] ?? 0;
95+
outputArray[i + 1] = compositeArray[i + 1] ?? 0;
96+
outputArray[i + 2] = compositeArray[i + 2] ?? 0;
97+
outputArray[i + 3] = Math.round((compositeArray[i + 3] ?? 0) * maskAlpha);
98+
}
99+
100+
// Package the masked pixels into an ImageData and draw them to an offscreen canvas.
101+
const outputImageData = new ImageData(outputArray, rect.width, rect.height);
102+
const outputCanvas = document.createElement('canvas');
103+
outputCanvas.width = rect.width;
104+
outputCanvas.height = rect.height;
105+
const outputContext = outputCanvas.getContext('2d');
106+
107+
if (!outputContext) {
108+
throw new Error('Failed to create canvas context for masked extraction');
109+
}
110+
111+
outputContext.putImageData(outputImageData, 0, 0);
112+
113+
// Upload the extracted canvas region as a real image resource and returns image_name
114+
115+
const blob = await canvasToBlob(outputCanvas);
116+
117+
const imageDTO = await uploadImage({
118+
file: new File([blob], 'inpaint-extract.png', { type: 'image/png' }),
119+
image_category: 'general',
120+
is_intermediate: true,
121+
silent: true,
122+
});
123+
124+
// Convert the uploaded image DTO into the canvas image state to avoid serializing the PNG in client state.
125+
const imageState: CanvasImageState = imageDTOToImageObject(imageDTO);
126+
127+
// Insert the new raster layer so it appears at the top of raster layer goup.
128+
const addAfter = canvasManager.stateApi.getRasterLayersState().entities.at(-1)?.id;
129+
130+
canvasManager.stateApi.addRasterLayer({
131+
overrides: {
132+
objects: [imageState],
133+
position: { x: rect.x, y: rect.y },
134+
},
135+
isSelected: true,
136+
addAfter,
137+
});
138+
} catch (error) {
139+
log.error({ error: serializeError(error as Error) }, 'Failed to extract masked area to raster layer');
140+
toast({ status: 'error', title: 'Unable to extract masked area.' });
141+
}
142+
})();
143+
}, [canvasManager, entityIdentifier]);
144+
145+
return (
146+
<MenuItem onClick={onExtract} icon={<PiSelectionBackgroundBold />} isDisabled={isBusy}>
147+
{t('controlLayers.extractRegion')}
148+
</MenuItem>
149+
);
150+
});
151+
152+
InpaintMaskMenuItemsExtractMaskedArea.displayName = 'InpaintMaskMenuItemsExtractMaskedArea';

0 commit comments

Comments
 (0)