Skip to content

Commit

Permalink
add the dead pill stop.txt
Browse files Browse the repository at this point in the history
  • Loading branch information
danielgatis committed Mar 12, 2024
1 parent b601763 commit d63da2e
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 3 deletions.
4 changes: 4 additions & 0 deletions rembg/bg.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))


Expand Down
30 changes: 27 additions & 3 deletions rembg/commands/p_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -115,6 +123,7 @@ def p_command(
input: pathlib.Path,
output: pathlib.Path,
watch: bool,
delete_input: bool,
**kwargs,
) -> None:
"""
Expand All @@ -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:
Expand Down Expand Up @@ -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)

Expand All @@ -179,21 +193,31 @@ 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()
observer.schedule(event_handler, input, recursive=False)
observer.start()

try:
while True:
while should_watch:
time.sleep(1)

finally:
Expand Down

0 comments on commit d63da2e

Please sign in to comment.