From 93b185dc3b27c943aa25e99ada1d6f55a8a58f5d Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Wed, 12 Jun 2024 17:12:53 +1000
Subject: [PATCH] feat(ui): bbox tool

---
 .../controlLayers/components/ToolChooser.tsx  | 21 ++++++++++++++++++-
 .../konva/renderers/previewLayer.ts           | 18 ++++++++--------
 .../src/features/controlLayers/store/types.ts |  2 +-
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/invokeai/frontend/web/src/features/controlLayers/components/ToolChooser.tsx b/invokeai/frontend/web/src/features/controlLayers/components/ToolChooser.tsx
index 0327d7de9b..9e26fed592 100644
--- a/invokeai/frontend/web/src/features/controlLayers/components/ToolChooser.tsx
+++ b/invokeai/frontend/web/src/features/controlLayers/components/ToolChooser.tsx
@@ -11,7 +11,14 @@ import {
 import { useCallback } from 'react';
 import { useHotkeys } from 'react-hotkeys-hook';
 import { useTranslation } from 'react-i18next';
-import { PiArrowsOutCardinalBold, PiEraserBold, PiHandBold, PiPaintBrushBold, PiRectangleBold } from 'react-icons/pi';
+import {
+  PiArrowsOutCardinalBold,
+  PiBoundingBoxBold,
+  PiEraserBold,
+  PiHandBold,
+  PiPaintBrushBold,
+  PiRectangleBold,
+} from 'react-icons/pi';
 
 const selectIsDisabled = createSelector(selectControlLayersSlice, (controlLayers) => {
   const selectedLayer = controlLayers.present.layers.find((l) => l.id === controlLayers.present.selectedLayerId);
@@ -45,6 +52,10 @@ export const ToolChooser: React.FC = () => {
     $tool.set('view');
   }, []);
   useHotkeys('h', setToolToView, { enabled: !isDisabled }, [isDisabled]);
+  const setToolToBbox = useCallback(() => {
+    $tool.set('bbox');
+  }, []);
+  useHotkeys('q', setToolToBbox, { enabled: !isDisabled }, [isDisabled]);
 
   const resetSelectedLayer = useCallback(() => {
     if (selectedLayerId === null) {
@@ -101,6 +112,14 @@ export const ToolChooser: React.FC = () => {
         onClick={setToolToView}
         isDisabled={isDisabled}
       />
+      <IconButton
+        aria-label={`${t('controlLayers.bbox')} (Q)`}
+        tooltip={`${t('controlLayers.bbox')} (Q)`}
+        icon={<PiBoundingBoxBold />}
+        variant={tool === 'bbox' ? 'solid' : 'outline'}
+        onClick={setToolToBbox}
+        isDisabled={isDisabled}
+      />
     </ButtonGroup>
   );
 };
diff --git a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/previewLayer.ts b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/previewLayer.ts
index d26f2a29cb..1947d868e3 100644
--- a/invokeai/frontend/web/src/features/controlLayers/konva/renderers/previewLayer.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/konva/renderers/previewLayer.ts
@@ -62,10 +62,10 @@ export const getBboxPreviewGroup = (
 
   // Use a transformer for the generation bbox. Transformers need some shape to transform, we will use a fully
   // transparent rect for this purpose.
-  bboxPreviewGroup = new Konva.Group({ id: PREVIEW_GENERATION_BBOX_GROUP });
+  bboxPreviewGroup = new Konva.Group({ id: PREVIEW_GENERATION_BBOX_GROUP, listening: false });
   const bboxRect = new Konva.Rect({
     id: PREVIEW_GENERATION_BBOX_DUMMY_RECT,
-    listening: true,
+    listening: false,
     strokeEnabled: false,
     draggable: true,
     ...getBbox(),
@@ -269,12 +269,10 @@ export const renderBboxPreview = (
   );
   const bboxRect = bboxGroup.findOne<Konva.Rect>(`#${PREVIEW_GENERATION_BBOX_DUMMY_RECT}`);
   const bboxTransformer = bboxGroup.findOne<Konva.Transformer>(`#${PREVIEW_GENERATION_BBOX_TRANSFORMER}`);
+  bboxGroup.listening(tool === 'bbox');
   // This updates the bbox during transformation
-  bboxRect?.setAttrs({ ...bbox, scaleX: 1, scaleY: 1, listening: tool === 'move' });
-  bboxTransformer?.setAttrs({
-    listening: tool === 'move',
-    enabledAnchors: tool === 'move' ? ALL_ANCHORS : NO_ANCHORS,
-  });
+  bboxRect?.setAttrs({ ...bbox, scaleX: 1, scaleY: 1, listening: tool === 'bbox' });
+  bboxTransformer?.setAttrs({ listening: tool === 'bbox', enabledAnchors: tool === 'bbox' ? ALL_ANCHORS : NO_ANCHORS });
 };
 
 export const getToolPreviewGroup = (stage: Konva.Stage): Konva.Group => {
@@ -365,9 +363,11 @@ export const renderToolPreview = (
   } else if (tool === 'rect') {
     // Rect gets a crosshair
     stage.container().style.cursor = 'crosshair';
-  } else {
-    // Else we hide the native cursor and use the konva-rendered brush preview
+  } else if (tool === 'brush' || tool === 'eraser') {
+    // Hide the native cursor and use the konva-rendered brush preview
     stage.container().style.cursor = 'none';
+  } else if (tool === 'bbox') {
+    stage.container().style.cursor = 'default';
   }
 
   stage.draggable(tool === 'view');
diff --git a/invokeai/frontend/web/src/features/controlLayers/store/types.ts b/invokeai/frontend/web/src/features/controlLayers/store/types.ts
index 860c8b5586..5380fe79d5 100644
--- a/invokeai/frontend/web/src/features/controlLayers/store/types.ts
+++ b/invokeai/frontend/web/src/features/controlLayers/store/types.ts
@@ -23,7 +23,7 @@ import type { IRect } from 'konva/lib/types';
 import type { ImageDTO } from 'services/api/types';
 import { z } from 'zod';
 
-const zTool = z.enum(['brush', 'eraser', 'move', 'rect', 'view']);
+const zTool = z.enum(['brush', 'eraser', 'move', 'rect', 'view', 'bbox']);
 export type Tool = z.infer<typeof zTool>;
 const zDrawingTool = zTool.extract(['brush', 'eraser']);