auto save result image when --output-dir exists

lama_cleaner/app/src/App.tsx

@@ -8,12 +8,14 @@ import {
   enableFileManagerState,
   fileState,
   isDisableModelSwitchState,
+  isEnableAutoSavingState,
   toastState,
 } from './store/Atoms'
 import { keepGUIAlive } from './utils'
 import Header from './components/Header/Header'
 import useHotKey from './hooks/useHotkey'
 import {
+  getEnableAutoSaving,
   getEnableFileManager,
   getIsDisableModelSwitch,
   isDesktop,
@@ -34,6 +36,7 @@ function App() {
   const userInputImage = useInputImage()
   const setIsDisableModelSwitch = useSetRecoilState(isDisableModelSwitchState)
   const setEnableFileManager = useSetRecoilState(enableFileManagerState)
+  const setIsEnableAutoSavingState = useSetRecoilState(isEnableAutoSavingState)
 
   // Set Input Image
   useEffect(() => {
@@ -66,7 +69,17 @@ function App() {
       setEnableFileManager(isEnabled === 'true')
     }
     fetchData2()
-  }, [setEnableFileManager, setIsDisableModelSwitch])
+
+    const fetchData3 = async () => {
+      const isEnabled = await getEnableAutoSaving().then(res => res.text())
+      setIsEnableAutoSavingState(isEnabled === 'true')
+    }
+    fetchData3()
+  }, [
+    setEnableFileManager,
+    setIsDisableModelSwitch,
+    setIsEnableAutoSavingState,
+  ])
 
   // Dark Mode Hotkey
   useHotKey(

lama_cleaner/app/src/adapters/inpainting.ts

@@ -122,6 +122,12 @@ export function getEnableFileManager() {
   })
 }
 
+export function getEnableAutoSaving() {
+  return fetch(`${API_ENDPOINT}/is_enable_auto_saving`, {
+    method: 'GET',
+  })
+}
+
 export function switchModel(name: string) {
   const fd = new FormData()
   fd.append('name', name)

lama_cleaner/app/src/components/Editor/Editor.tsx

@@ -20,7 +20,6 @@ import { useRecoilState, useRecoilValue, useSetRecoilState } from 'recoil'
 import { useWindowSize, useKey, useKeyPressEvent } from 'react-use'
 import inpaint, {
   downloadToOutput,
-  makeGif,
   postInteractiveSeg,
 } from '../../adapters/inpainting'
 import Button from '../shared/Button'
@@ -45,12 +44,11 @@ import {
   imageWidthState,
   interactiveSegClicksState,
   isDiffusionModelsState,
+  isEnableAutoSavingState,
   isInpaintingState,
   isInteractiveSegRunningState,
   isInteractiveSegState,
-  isPaintByExampleState,
   isPix2PixState,
-  isSDState,
   negativePropmtState,
   propmtState,
   runManuallyState,
@@ -184,6 +182,7 @@ export default function Editor() {
   const [redoCurLines, setRedoCurLines] = useState<Line[]>([])
   const [redoLineGroups, setRedoLineGroups] = useState<LineGroup[]>([])
   const enableFileManager = useRecoilValue(enableFileManagerState)
+  const isEnableAutoSaving = useRecoilValue(isEnableAutoSavingState)
 
   const setImageWidth = useSetRecoilState(imageWidthState)
   const setImageHeight = useSetRecoilState(imageHeightState)
@@ -1101,7 +1100,7 @@ export default function Editor() {
     if (file === undefined) {
       return
     }
-    if (enableFileManager && renders.length > 0) {
+    if ((enableFileManager || isEnableAutoSaving) && renders.length > 0) {
       try {
         downloadToOutput(renders[renders.length - 1], file.name, file.type)
         setToastState({

lama_cleaner/app/src/store/Atoms.tsx

@@ -41,6 +41,7 @@ interface AppState {
   disableShortCuts: boolean
   isInpainting: boolean
   isDisableModelSwitch: boolean
+  isEnableAutoSaving: boolean
   isInteractiveSeg: boolean
   isInteractiveSegRunning: boolean
   interactiveSegClicks: number[][]
@@ -58,6 +59,7 @@ export const appState = atom<AppState>({
     disableShortCuts: false,
     isInpainting: false,
     isDisableModelSwitch: false,
+    isEnableAutoSaving: false,
     isInteractiveSeg: false,
     isInteractiveSegRunning: false,
     interactiveSegClicks: [],
@@ -221,6 +223,18 @@ export const isDisableModelSwitchState = selector({
   },
 })
 
+export const isEnableAutoSavingState = selector({
+  key: 'isEnableAutoSavingState',
+  get: ({ get }) => {
+    const app = get(appState)
+    return app.isEnableAutoSaving
+  },
+  set: ({ get, set }, newValue: any) => {
+    const app = get(appState)
+    set(appState, { ...app, isEnableAutoSaving: newValue })
+  },
+})
+
 export const croperState = atom<Rect>({
   key: 'croperState',
   default: {

lama_cleaner/const.py

@@ -59,7 +59,7 @@ Model download directory (by setting XDG_CACHE_HOME environment variable), by de
 """
 
 OUTPUT_DIR_HELP = """
-Only required when --input is directory. Result images will be saved to output directory automatically.
+Result images will be saved to output directory automatically without confirmation.
 """
 
 INPUT_HELP = """

lama_cleaner/file_manager/file_manager.py

@@ -14,7 +14,7 @@ from PIL import Image, ImageOps, PngImagePlugin
 from loguru import logger
 
 LARGE_ENOUGH_NUMBER = 100
-PngImagePlugin.MAX_TEXT_CHUNK = LARGE_ENOUGH_NUMBER * (1024 ** 2)
+PngImagePlugin.MAX_TEXT_CHUNK = LARGE_ENOUGH_NUMBER * (1024**2)
 from .storage_backends import FilesystemStorageBackend
 from .utils import aspect_to_string, generate_filename, glob_img
 
@@ -63,11 +63,11 @@ class FileManager(FileSystemEventHandler):
         if event.src_path == str(self.root_directory):
             logger.info(f"Image directory {event.src_path} modified")
             self.image_dir_filenames = self._media_names(self.root_directory)
-            self.modified_time['image'] = datetime.utcnow()
+            self.modified_time["image"] = datetime.utcnow()
         elif event.src_path == str(self.output_dir):
             logger.info(f"Output directory {event.src_path} modified")
             self.output_dir_filenames = self._media_names(self.output_dir)
-            self.modified_time['output'] = datetime.utcnow()
+            self.modified_time["output"] = datetime.utcnow()
 
     def init_app(self, app):
         if self.app is None:
@@ -83,21 +83,15 @@ class FileManager(FileSystemEventHandler):
         app.extensions["thumbnail"] = self
 
         app.config.setdefault("THUMBNAIL_MEDIA_ROOT", self._default_root_directory)
-        app.config.setdefault("THUMBNAIL_MEDIA_THUMBNAIL_ROOT", self._default_thumbnail_directory)
+        app.config.setdefault(
+            "THUMBNAIL_MEDIA_THUMBNAIL_ROOT", self._default_thumbnail_directory
+        )
         app.config.setdefault("THUMBNAIL_MEDIA_URL", self._default_root_url)
-        app.config.setdefault("THUMBNAIL_MEDIA_THUMBNAIL_URL", self._default_thumbnail_root_url)
+        app.config.setdefault(
+            "THUMBNAIL_MEDIA_THUMBNAIL_URL", self._default_thumbnail_root_url
+        )
         app.config.setdefault("THUMBNAIL_DEFAULT_FORMAT", self._default_format)
 
-    def save_to_output_directory(self, image: np.ndarray, filename: str):
-        fp = Path(filename)
-        new_name = fp.stem + f"_{int(time.time())}" + fp.suffix
-        if image.shape[2] == 3:
-            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-        elif image.shape[2] == 4:
-            image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGRA)
-
-        cv2.imwrite(str(self.output_dir / new_name), image)
-
     @property
     def root_directory(self):
         path = self.app.config["THUMBNAIL_MEDIA_ROOT"]
@@ -137,14 +131,23 @@ class FileManager(FileSystemEventHandler):
         for name in names:
             path = os.path.join(directory, name)
             img = Image.open(path)
-            res.append({"name": name, "height": img.height, "width": img.width, "ctime": os.path.getctime(path)})
+            res.append(
+                {
+                    "name": name,
+                    "height": img.height,
+                    "width": img.width,
+                    "ctime": os.path.getctime(path),
+                }
+            )
         return res
 
     @property
     def thumbnail_url(self):
         return self.app.config["THUMBNAIL_MEDIA_THUMBNAIL_URL"]
 
-    def get_thumbnail(self, directory: Path, original_filename: str, width, height, **options):
+    def get_thumbnail(
+        self, directory: Path, original_filename: str, width, height, **options
+    ):
         storage = FilesystemStorageBackend(self.app)
         crop = options.get("crop", "fit")
         background = options.get("background")
@@ -163,13 +166,19 @@ class FileManager(FileSystemEventHandler):
         thumbnail_size = (width, height)
 
         thumbnail_filename = generate_filename(
-            original_filename, aspect_to_string(thumbnail_size), crop, background, quality
+            original_filename,
+            aspect_to_string(thumbnail_size),
+            crop,
+            background,
+            quality,
         )
 
         thumbnail_filepath = os.path.join(
             self.thumbnail_directory, original_path, thumbnail_filename
         )
-        thumbnail_url = os.path.join(self.thumbnail_url, original_path, thumbnail_filename)
+        thumbnail_url = os.path.join(
+            self.thumbnail_url, original_path, thumbnail_filename
+        )
 
         if storage.exists(thumbnail_filepath):
             return thumbnail_url, (width, height)
@@ -183,7 +192,9 @@ class FileManager(FileSystemEventHandler):
         # get original image format
         options["format"] = options.get("format", image.format)
 
-        image = self._create_thumbnail(image, thumbnail_size, crop, background=background)
+        image = self._create_thumbnail(
+            image, thumbnail_size, crop, background=background
+        )
 
         raw_data = self.get_raw_data(image, **options)
         storage.save(thumbnail_filepath, raw_data)

lama_cleaner/parse_args.py

@@ -170,15 +170,14 @@ def parse_args():
                 parser.error(
                     f"invalid --input: {args.input} is a directory, --output-dir is required"
                 )
-            else:
+
-                output_dir = Path(args.output_dir)
+    if args.output_dir is not None:
-                if not output_dir.exists():
+        output_dir = Path(args.output_dir)
-                    logger.info(f"Creating output directory: {output_dir}")
+        if not output_dir.exists():
-                    output_dir.mkdir(parents=True)
+            logger.info(f"Creating output directory: {output_dir}")
-                else:
+            output_dir.mkdir(parents=True)
-                    if not output_dir.is_dir():
+        else:
-                        parser.error(
+            if not output_dir.is_dir():
-                            f"invalid --output-dir: {output_dir} is not a directory"
+                parser.error(f"invalid --output-dir: {output_dir} is not a directory")
-                        )
 
     return args

lama_cleaner/server.py

@@ -83,11 +83,13 @@ CORS(app, expose_headers=["Content-Disposition"])
 
 model: ModelManager = None
 thumb: FileManager = None
+output_dir: str = None
 interactive_seg_model: InteractiveSeg = None
 device = None
 input_image_path: str = None
 is_disable_model_switch: bool = False
 is_enable_file_manager: bool = False
+is_enable_auto_saving: bool = False
 is_desktop: bool = False
 
 
@@ -124,11 +126,20 @@ def make_gif():
 
 @app.route("/save_image", methods=["POST"])
 def save_image():
-    # all image in output directory
+    if output_dir is None:
+        return "--output-dir is None", 500
+
     input = request.files
+    filename = request.form["filename"]
     origin_image_bytes = input["image"].read()  # RGB
     image, _ = load_img(origin_image_bytes)
-    thumb.save_to_output_directory(image, request.form["filename"])
+    if image.shape[2] == 3:
+        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    elif image.shape[2] == 4:
+        image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGRA)
+
+    cv2.imwrite(os.path.join(output_dir, filename), image)
+
     return "ok", 200
 
 
@@ -348,6 +359,12 @@ def get_is_enable_file_manager():
     return res, 200
 
 
+@app.route("/is_enable_auto_saving")
+def get_is_enable_auto_saving():
+    res = "true" if is_enable_auto_saving else "false"
+    return res, 200
+
+
 @app.route("/model_downloaded/<name>")
 def model_downloaded(name):
     return str(model.is_downloaded(name)), 200
@@ -408,6 +425,12 @@ def main(args):
     global is_enable_file_manager
     global is_desktop
     global thumb
+    global output_dir
+    global is_enable_auto_saving
+
+    output_dir = args.output_dir
+    if output_dir is not None:
+        is_enable_auto_saving = True
 
     device = torch.device(args.device)
     is_disable_model_switch = args.disable_model_switch