From 1f8fb295945ba805991dbd0a04448ac99b7c7541 Mon Sep 17 00:00:00 2001
From: Qing <cwq1913@gmail.com>
Date: Tue, 12 Jul 2022 22:26:28 +0800
Subject: [PATCH] add Download mask settings

---
 .../app/src/components/Editor/Editor.tsx      | 15 ++++++++++
 .../Settings/DownloadMaskSettingBlock.tsx     | 29 +++++++++++++++++++
 .../src/components/Settings/SettingBlock.tsx  |  2 +-
 .../src/components/Settings/SettingsModal.tsx |  2 ++
 lama_cleaner/app/src/store/Atoms.tsx          |  2 ++
 5 files changed, 49 insertions(+), 1 deletion(-)
 create mode 100644 lama_cleaner/app/src/components/Settings/DownloadMaskSettingBlock.tsx

diff --git a/lama_cleaner/app/src/components/Editor/Editor.tsx b/lama_cleaner/app/src/components/Editor/Editor.tsx
index 156ca78..63cc0bc 100644
--- a/lama_cleaner/app/src/components/Editor/Editor.tsx
+++ b/lama_cleaner/app/src/components/Editor/Editor.tsx
@@ -22,6 +22,7 @@ import Button from '../shared/Button'
 import Slider from './Slider'
 import SizeSelector from './SizeSelector'
 import {
+  dataURItoBlob,
   downloadImage,
   isMidClick,
   isRightClick,
@@ -648,6 +649,20 @@ export default function Editor(props: EditorProps) {
     const name = file.name.replace(/(\.[\w\d_-]+)$/i, '_cleanup$1')
     const curRender = renders[renders.length - 1]
     downloadImage(curRender.currentSrc, name)
+    if (settings.downloadMask) {
+      let maskFileName = file.name.replace(/(\.[\w\d_-]+)$/i, '_mask$1')
+      maskFileName = maskFileName.replace(/\.[^/.]+$/, '.jpg')
+
+      drawLinesOnMask(lineGroups)
+      // Create a link
+      const aDownloadLink = document.createElement('a')
+      // Add the name of the file to the link
+      aDownloadLink.download = maskFileName
+      // Attach the data to the link
+      aDownloadLink.href = maskCanvas.toDataURL('image/jpeg')
+      // Get the code to click the download link
+      aDownloadLink.click()
+    }
   }
 
   const onSizeLimitChange = (_sizeLimit: number) => {
diff --git a/lama_cleaner/app/src/components/Settings/DownloadMaskSettingBlock.tsx b/lama_cleaner/app/src/components/Settings/DownloadMaskSettingBlock.tsx
new file mode 100644
index 0000000..985d23d
--- /dev/null
+++ b/lama_cleaner/app/src/components/Settings/DownloadMaskSettingBlock.tsx
@@ -0,0 +1,29 @@
+import React from 'react'
+import { useRecoilState } from 'recoil'
+import { settingState } from '../../store/Atoms'
+import { Switch, SwitchThumb } from '../shared/Switch'
+import SettingBlock from './SettingBlock'
+
+const DownloadMaskSettingBlock: React.FC = () => {
+  const [setting, setSettingState] = useRecoilState(settingState)
+
+  const onCheckChange = (checked: boolean) => {
+    setSettingState(old => {
+      return { ...old, downloadMask: checked }
+    })
+  }
+
+  return (
+    <SettingBlock
+      title="Download Mask"
+      desc="Download inpainting result and mask"
+      input={
+        <Switch checked={setting.downloadMask} onCheckedChange={onCheckChange}>
+          <SwitchThumb />
+        </Switch>
+      }
+    />
+  )
+}
+
+export default DownloadMaskSettingBlock
diff --git a/lama_cleaner/app/src/components/Settings/SettingBlock.tsx b/lama_cleaner/app/src/components/Settings/SettingBlock.tsx
index 026cde3..a984183 100644
--- a/lama_cleaner/app/src/components/Settings/SettingBlock.tsx
+++ b/lama_cleaner/app/src/components/Settings/SettingBlock.tsx
@@ -17,7 +17,7 @@ function SettingBlock(props: SettingBlockProps) {
         <div className="setting-block-content-title">
           <span>{title}</span>
           {desc && (
-            <Tooltip content={<div style={{ width: 400 }}>{desc}</div>}>
+            <Tooltip content={<div style={{ maxWidth: 400 }}>{desc}</div>}>
               <svg
                 width="18"
                 height="18"
diff --git a/lama_cleaner/app/src/components/Settings/SettingsModal.tsx b/lama_cleaner/app/src/components/Settings/SettingsModal.tsx
index 06d7033..0d441cb 100644
--- a/lama_cleaner/app/src/components/Settings/SettingsModal.tsx
+++ b/lama_cleaner/app/src/components/Settings/SettingsModal.tsx
@@ -7,6 +7,7 @@ import ManualRunInpaintingSettingBlock from './ManualRunInpaintingSettingBlock'
 import HDSettingBlock from './HDSettingBlock'
 import ModelSettingBlock from './ModelSettingBlock'
 import GraduallyInpaintingSettingBlock from './GraduallyInpaintingSettingBlock'
+import DownloadMaskSettingBlock from './DownloadMaskSettingBlock'
 
 interface SettingModalProps {
   onClose: () => void
@@ -31,6 +32,7 @@ export default function SettingModal(props: SettingModalProps) {
     >
       <ManualRunInpaintingSettingBlock />
       <GraduallyInpaintingSettingBlock />
+      <DownloadMaskSettingBlock />
       <ModelSettingBlock />
       <HDSettingBlock />
     </Modal>
diff --git a/lama_cleaner/app/src/store/Atoms.tsx b/lama_cleaner/app/src/store/Atoms.tsx
index 2b73cb3..89ab854 100644
--- a/lama_cleaner/app/src/store/Atoms.tsx
+++ b/lama_cleaner/app/src/store/Atoms.tsx
@@ -32,6 +32,7 @@ export const shortcutsState = atom<boolean>({
 
 export interface Settings {
   show: boolean
+  downloadMask: boolean
   graduallyInpainting: boolean
   runInpaintingManually: boolean
   model: AIModel
@@ -49,6 +50,7 @@ export interface Settings {
 
 export const settingStateDefault = {
   show: false,
+  downloadMask: false,
   graduallyInpainting: true,
   runInpaintingManually: false,
   model: AIModel.LAMA,