From d63da2e6fbca07686166fbb7bbe3bf8f38d808a1 Mon Sep 17 00:00:00 2001
From: Daniel Gatis <danielgatis@gmail.com>
Date: Tue, 12 Mar 2024 19:43:25 -0300
Subject: [PATCH] add the dead pill stop.txt

---
 rembg/bg.py                 |  4 ++++
 rembg/commands/p_command.py | 30 +++++++++++++++++++++++++++---
 2 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/rembg/bg.py b/rembg/bg.py
index 8629f827..170329cd 100644
--- a/rembg/bg.py
+++ b/rembg/bg.py
@@ -22,6 +22,10 @@
 from .sessions import sessions_class
 from .sessions.base import BaseSession
 
+import onnxruntime as ort
+
+ort.set_default_logger_severity(3)
+
 kernel = getStructuringElement(MORPH_ELLIPSE, (3, 3))
 
 
diff --git a/rembg/commands/p_command.py b/rembg/commands/p_command.py
index 41c81467..4f300a00 100644
--- a/rembg/commands/p_command.py
+++ b/rembg/commands/p_command.py
@@ -80,6 +80,14 @@
     show_default=True,
     help="watches a folder for changes",
 )
+@click.option(
+    "-d",
+    "--delete_input",
+    default=False,
+    is_flag=True,
+    show_default=True,
+    help="delete input file after processing",
+)
 @click.option(
     "-bgc",
     "--bgcolor",
@@ -115,6 +123,7 @@ def p_command(
     input: pathlib.Path,
     output: pathlib.Path,
     watch: bool,
+    delete_input: bool,
     **kwargs,
 ) -> None:
     """
@@ -132,6 +141,7 @@ def p_command(
         input (pathlib.Path): The path to the input folder.
         output (pathlib.Path): The path to the output folder.
         watch (bool): Whether to watch the input folder for changes.
+        delete_input (bool): Whether to delete the input file after processing.
         **kwargs: Additional keyword arguments.
 
     Returns:
@@ -167,6 +177,10 @@ def process(each_input: pathlib.Path) -> None:
                     print(
                         f"processed: {each_input.absolute()} -> {each_output.absolute()}"
                     )
+
+            if delete_input:
+                each_input.unlink()
+
         except Exception as e:
             print(e)
 
@@ -179,13 +193,23 @@ def process(each_input: pathlib.Path) -> None:
             process(each_input)
 
     if watch:
+        should_watch = True
         observer = Observer()
 
         class EventHandler(FileSystemEventHandler):
             def on_any_event(self, event: FileSystemEvent) -> None:
-                if not (
-                    event.is_directory or event.event_type in ["deleted", "closed"]
+                if (
+                    not (
+                        event.is_directory or event.event_type in ["deleted", "closed"]
+                    )
+                    and pathlib.Path(event.src_path).exists()
                 ):
+                    if event.src_path.endswith("stop.txt"):
+                        nonlocal should_watch
+                        should_watch = False
+                        pathlib.Path(event.src_path).unlink()
+                        return
+
                     process(pathlib.Path(event.src_path))
 
         event_handler = EventHandler()
@@ -193,7 +217,7 @@ def on_any_event(self, event: FileSystemEvent) -> None:
         observer.start()
 
         try:
-            while True:
+            while should_watch:
                 time.sleep(1)
 
         finally: