From f039b381c9f350579dce2ee6d0ec9bc1dd9e5cf6 Mon Sep 17 00:00:00 2001
From: Abdullah Alfaraj <7842232+AbdullahAlfaraj@users.noreply.github.com>
Date: Tue, 21 Nov 2023 05:34:27 +0300
Subject: [PATCH] Added option in lasso inpaint to maintain original selection
 ratio instead of squaring.

---
 selection.js                  | 36 ++++++++++++++++++++++++++++++-----
 typescripts/sd_tab/sd_tab.tsx | 16 ++++++++++++++++
 typescripts/sd_tab/util.ts    |  1 +
 typescripts/session/modes.ts  |  3 ++-
 4 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/selection.js b/selection.js
index 412b606c..98782bcd 100644
--- a/selection.js
+++ b/selection.js
@@ -324,6 +324,24 @@ async function channelToSelectionExe(channel_name = 'mask') {
     }
 }
 
+function keepRatio(selectionInfo, offset) {
+    // Calculate the current width and height
+    let width = selectionInfo.right - selectionInfo.left
+    let height = selectionInfo.bottom - selectionInfo.top
+
+    // Calculate the new coordinates with the offset
+    selectionInfo.right += offset
+    selectionInfo.left -= offset
+    selectionInfo.bottom += offset
+    selectionInfo.top -= offset
+
+    // Update width and height
+    selectionInfo.width = width + 2 * offset
+    selectionInfo.height = height + 2 * offset
+
+    return selectionInfo
+}
+
 function makeSquare(selectionInfo, offset) {
     // Calculate the current width and height
     let width = selectionInfo.right - selectionInfo.left
@@ -349,12 +367,20 @@ function makeSquare(selectionInfo, offset) {
     return selectionInfo
 }
 
-async function inpaintLassoInitImageAndMask(channel_name = 'mask', offset = 0) {
+async function inpaintLassoInitImageAndMask(
+    channel_name = 'mask',
+    offset = 0,
+    make_square = true
+) {
     const selectionInfo = await psapi.getSelectionInfoExe()
     //convert the selection box into square box so that you have best output results
-    const squareSelection = makeSquare(selectionInfo, offset)
+
+    const newSelection = make_square
+        ? makeSquare(selectionInfo, offset)
+        : keepRatio(selectionInfo, offset)
+
     //correct width and height sliders, since this is lasso mode.
-    await calcWidthHeightFromSelection(squareSelection)
+    await calcWidthHeightFromSelection(newSelection)
 
     async function getImageFromCanvas() {
         const width = html_manip.getWidth()
@@ -363,7 +389,7 @@ async function inpaintLassoInitImageAndMask(channel_name = 'mask', offset = 0) {
         const base64 = await io.IO.getSelectionFromCanvasAsBase64Interface_New(
             width,
             height,
-            squareSelection,
+            newSelection,
             true
         )
         return base64
@@ -396,7 +422,7 @@ async function inpaintLassoInitImageAndMask(channel_name = 'mask', offset = 0) {
             synchronousExecution: true,
         })
         // const selection_info = await psapi.getSelectionInfoExe()
-        mask_base64 = await fillSelectionWhiteOutsideBlack(squareSelection)
+        mask_base64 = await fillSelectionWhiteOutsideBlack(newSelection)
     })
 
     //save laso selection to channel
diff --git a/typescripts/sd_tab/sd_tab.tsx b/typescripts/sd_tab/sd_tab.tsx
index 72b365b8..5890f8af 100644
--- a/typescripts/sd_tab/sd_tab.tsx
+++ b/typescripts/sd_tab/sd_tab.tsx
@@ -105,6 +105,22 @@ const Modes = observer(() => {
                         >
                             Lasso Mode
                         </SpCheckBox>
+                        <SpCheckBox
+                            style={{
+                                marginLeft: '10px',
+                                display: store.data.is_lasso_mode
+                                    ? void 0
+                                    : 'none',
+                            }}
+                            onChange={() => {
+                                helper_store.data.make_square =
+                                    !helper_store.data.make_square
+                            }}
+                            checked={helper_store.data.make_square}
+                            // id={`chEnableControlNet_${this.props.index}`}
+                        >
+                            Make Square
+                        </SpCheckBox>
 
                         <SpSlider
                             show-value="false"
diff --git a/typescripts/sd_tab/util.ts b/typescripts/sd_tab/util.ts
index 95d87729..ce14ff79 100644
--- a/typescripts/sd_tab/util.ts
+++ b/typescripts/sd_tab/util.ts
@@ -206,6 +206,7 @@ export const helper_store = new AStore({
     native_presets: {},
     base_size: 512 as number,
     lasso_offset: 10 as number,
+    make_square: true as boolean,
 })
 export async function refreshModels() {
     let b_result = false
diff --git a/typescripts/session/modes.ts b/typescripts/session/modes.ts
index 082040d5..514e8413 100644
--- a/typescripts/session/modes.ts
+++ b/typescripts/session/modes.ts
@@ -588,7 +588,8 @@ export class LassoInpaintMode extends Img2ImgMode {
         }
         const [init_image, mask] = await selection.inpaintLassoInitImageAndMask(
             'mask',
-            sd_tab_util.helper_store.data.lasso_offset
+            sd_tab_util.helper_store.data.lasso_offset,
+            sd_tab_util.helper_store.data.make_square
         )
 
         const selectionInfo = await psapi.getSelectionInfoExe()