diff --git a/lama_cleaner/app/src/components/FileManager/FileManager.tsx b/lama_cleaner/app/src/components/FileManager/FileManager.tsx index bae6618..6cf2885 100644 --- a/lama_cleaner/app/src/components/FileManager/FileManager.tsx +++ b/lama_cleaner/app/src/components/FileManager/FileManager.tsx @@ -109,6 +109,9 @@ export default function FileManager(props: Props) { ) useEffect(() => { + if (!show) { + return + } const fetchData = async () => { try { const filenames = await getMedias(tab) @@ -141,7 +144,15 @@ export default function FileManager(props: Props) { } } fetchData() - }, [setToastState, tab, debouncedSearchText, sortBy, sortOrder, photoWidth]) + }, [ + setToastState, + tab, + debouncedSearchText, + sortBy, + sortOrder, + photoWidth, + show, + ]) const onScroll = (event: SyntheticEvent) => { setScrollTop(event.currentTarget.scrollTop) diff --git a/lama_cleaner/file_manager/file_manager.py b/lama_cleaner/file_manager/file_manager.py index 2c46669..9e4cfb4 100644 --- a/lama_cleaner/file_manager/file_manager.py +++ b/lama_cleaner/file_manager/file_manager.py @@ -1,13 +1,17 @@ # Copy from https://github.com/silentsokolov/flask-thumbnails/blob/master/flask_thumbnails/thumbnail.py import os -from cachetools import TTLCache, cached +from datetime import datetime + import cv2 import time from io import BytesIO from pathlib import Path import numpy as np +from watchdog.events import FileSystemEventHandler +from watchdog.observers import Observer from PIL import Image, ImageOps, PngImagePlugin +from loguru import logger LARGE_ENOUGH_NUMBER = 100 PngImagePlugin.MAX_TEXT_CHUNK = LARGE_ENOUGH_NUMBER * (1024 ** 2) @@ -15,7 +19,7 @@ from .storage_backends import FilesystemStorageBackend from .utils import aspect_to_string, generate_filename, glob_img -class FileManager: +class FileManager(FileSystemEventHandler): def __init__(self, app=None): self.app = app self._default_root_directory = "media" @@ -28,6 +32,43 @@ class FileManager: if app is not None: self.init_app(app) + self.image_dir_filenames = [] + self.output_dir_filenames = [] + + self.image_dir_observer = None + self.output_dir_observer = None + + self.modified_time = { + "image": datetime.utcnow(), + "output": datetime.utcnow(), + } + + def start(self): + self.image_dir_filenames = self._media_names(self.root_directory) + self.output_dir_filenames = self._media_names(self.output_dir) + + logger.info(f"Start watching image directory: {self.root_directory}") + self.image_dir_observer = Observer() + self.image_dir_observer.schedule(self, self.root_directory, recursive=False) + self.image_dir_observer.start() + + logger.info(f"Start watching output directory: {self.output_dir}") + self.output_dir_observer = Observer() + self.output_dir_observer.schedule(self, self.output_dir, recursive=False) + self.output_dir_observer.start() + + def on_modified(self, event): + if not os.path.isdir(event.src_path): + return + 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() + 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() + def init_app(self, app): if self.app is None: self.app = app @@ -80,14 +121,12 @@ class FileManager: return self.app.config["THUMBNAIL_MEDIA_URL"] @property - @cached(cache=TTLCache(maxsize=1024, ttl=30)) def media_names(self): - return self._media_names(self.root_directory) + return self.image_dir_filenames @property - @cached(cache=TTLCache(maxsize=1024, ttl=30)) def output_media_names(self): - return self._media_names(self.output_dir) + return self.output_dir_filenames @staticmethod def _media_names(directory: Path): diff --git a/lama_cleaner/server.py b/lama_cleaner/server.py index 16200ba..4345842 100644 --- a/lama_cleaner/server.py +++ b/lama_cleaner/server.py @@ -16,6 +16,7 @@ import cv2 import torch import numpy as np from loguru import logger +from watchdog.events import FileSystemEventHandler from lama_cleaner.interactive_seg import InteractiveSeg, Click from lama_cleaner.model_manager import ModelManager @@ -71,7 +72,7 @@ app.config["JSON_AS_ASCII"] = False CORS(app, expose_headers=["Content-Disposition"]) model: ModelManager = None -thumb = FileManager(app) +thumb: FileManager = None interactive_seg_model: InteractiveSeg = None device = None input_image_path: str = None @@ -105,9 +106,14 @@ def save_image(): @app.route("/medias/") def medias(tab): if tab == 'image': - # all images in input folder - return jsonify(thumb.media_names), 200 - return jsonify(thumb.output_media_names), 200 + response = make_response(jsonify(thumb.media_names), 200) + else: + response = make_response(jsonify(thumb.output_media_names), 200) + response.last_modified = thumb.modified_time[tab] + # response.cache_control.no_cache = True + # response.cache_control.max_age = 0 + response.make_conditional(request) + return response @app.route('/media//') @@ -340,6 +346,11 @@ def set_input_photo(): return "No Input Image" +class FSHandler(FileSystemEventHandler): + def on_modified(self, event): + print("File modified: %s" % event.src_path) + + def main(args): global model global interactive_seg_model @@ -348,6 +359,7 @@ def main(args): global is_disable_model_switch global is_enable_file_manager global is_desktop + global thumb device = torch.device(args.device) is_disable_model_switch = args.disable_model_switch @@ -356,10 +368,22 @@ def main(args): logger.info(f"Start with --disable-model-switch, model switch on frontend is disable") if os.path.isdir(args.input): + logger.info(f"Initialize file manager") + thumb = FileManager(app) + is_enable_file_manager = True app.config["THUMBNAIL_MEDIA_ROOT"] = args.input app.config["THUMBNAIL_MEDIA_THUMBNAIL_ROOT"] = os.path.join(args.output_dir, 'lama_cleaner_thumbnails') - is_enable_file_manager = True thumb.output_dir = Path(args.output_dir) + thumb.start() + # try: + # while True: + # time.sleep(1) + # finally: + # thumb.image_dir_observer.stop() + # thumb.image_dir_observer.join() + # thumb.output_dir_observer.stop() + # thumb.output_dir_observer.join() + else: input_image_path = args.input @@ -388,5 +412,4 @@ def main(args): ) ui.run() else: - # TODO: socketio app.run(host=args.host, port=args.port, debug=args.debug) diff --git a/requirements.txt b/requirements.txt index 36115ae..1b20495 100644 --- a/requirements.txt +++ b/requirements.txt @@ -12,4 +12,4 @@ markupsafe==2.0.1 scikit-image==0.19.3 diffusers[torch]==0.10.2 transformers>=4.25.1 -cachetools==5.2.0 \ No newline at end of file +watchdog==2.2.1 \ No newline at end of file