You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

daemon.py 3.3KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108
  1. """daemon logic."""
  2. import copy
  3. import os
  4. import sys
  5. import time
  6. from watchdog.events import FileSystemEventHandler
  7. from watchdog.observers import Observer
  8. from config import Config as Conf
  9. from processing.folder import FolderImageProcessing
  10. from transform.gan.mask import CorrectToMask, MaskrefToMaskdet, MaskfinToNude
  11. from transform.opencv.correct import DressToCorrect
  12. from transform.opencv.mask import MaskToMaskref, MaskdetToMaskfin
  13. class Watcher:
  14. """Watch a directory change."""
  15. def __init__(self, watching_dir, out_dir):
  16. """
  17. Watcher constructor.
  18. :param watching_dir: <string> directory to watch
  19. :param out_dir: <string> directory where save transformations
  20. """
  21. self.__observer = Observer()
  22. self.__watching_dir = watching_dir
  23. self.__out_dir = out_dir
  24. for k, v in {"Watched": self.__watching_dir, "Output": self.__out_dir}.items():
  25. Conf.log.info("{} Directory Is {}".format(k, v))
  26. if not os.path.isdir(self.__watching_dir):
  27. Conf.log.error("{} Directory {} Doesn't Exit.".format(k, v))
  28. sys.exit(0)
  29. def run(self):
  30. """
  31. Run the Watcher.
  32. :return: None
  33. """
  34. event_handler = Handler(self.__out_dir)
  35. self.__observer.schedule(event_handler, self.__watching_dir, recursive=True)
  36. self.__observer.start()
  37. try:
  38. while True:
  39. time.sleep(5)
  40. except KeyboardInterrupt:
  41. self.__observer.stop()
  42. except Exception as e:
  43. self.__observer.stop()
  44. Conf.log.error(e)
  45. Conf.log.error("An Unhandled Error Occurred.")
  46. sys.exit(1)
  47. self.__observer.join()
  48. class Handler(FileSystemEventHandler):
  49. """Handle a change in a watch directory."""
  50. def __init__(self, out_dir):
  51. """
  52. Create an Handler.
  53. :param out_dir: <string> directory where save transformations
  54. """
  55. self.__out_dir = out_dir
  56. self.__phases = [
  57. DressToCorrect, CorrectToMask, MaskToMaskref, MaskrefToMaskdet, MaskdetToMaskfin, MaskfinToNude
  58. ]
  59. def on_created(self, event):
  60. """
  61. Call when a file or directory is created.
  62. :param event: <DirCreatedEvent|FileCreatedEvent> trigger event
  63. :return: None
  64. """
  65. if not event.is_directory:
  66. Conf.log.debug("Received file created event {}.".format(event.src_path))
  67. if os.path.basename(event.src_path) == ".start":
  68. os.remove(event.src_path)
  69. start = time.time()
  70. Conf.log.info("Execution Of {} Folder.".format(os.path.dirname(event.src_path)))
  71. args = copy.deepcopy(Conf.args)
  72. args.update({
  73. "input": os.path.dirname(event.src_path),
  74. "output": self.__out_dir,
  75. })
  76. FolderImageProcessing(args=args).run()
  77. Conf.log.info("Execution of {} Folder Done in {}.".format(
  78. os.path.dirname(event.src_path), round(time.time() - start, 2)
  79. ))
  80. def main(_):
  81. """
  82. Start daemon main logic.
  83. :param _: None
  84. :return: None
  85. """
  86. Conf.log.info("Welcome to Dreampower Daemon")
  87. Watcher(Conf.args['input'], Conf.args['output']).run()