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.

main.py 4.5KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146
  1. """main logic."""
  2. import os
  3. import sys
  4. import time
  5. from multiprocessing import freeze_support
  6. import colorama
  7. import argv
  8. from config import Config as Conf
  9. from utils import check_shape
  10. from processing import SimpleTransform, FolderImageTransform, MultipleImageTransform
  11. from transform.gan.mask import CorrectToMask, MaskrefToMaskdet, MaskfinToNude
  12. from transform.opencv.resize import ImageToCrop, ImageToOverlay, ImageToRescale, ImageToResized, ImageToResizedCrop
  13. from transform.opencv.correct import DressToCorrect, ColorTransfer
  14. from transform.opencv.mask import MaskToMaskref, MaskdetToMaskfin
  15. def main(_):
  16. """Start main logic."""
  17. Conf.log.info("Welcome to DreamPower")
  18. if Conf.args['gpu_ids']:
  19. Conf.log.info("GAN Processing Will Use GPU IDs: {}".format(Conf.args['gpu_ids']))
  20. else:
  21. Conf.log.info("GAN Processing Will Use CPU")
  22. # Processing
  23. start = time.time()
  24. select_processing().run()
  25. Conf.log.info("Done! We have taken {} seconds".format(round(time.time() - start, 2)))
  26. # Exit
  27. sys.exit()
  28. def select_phases():
  29. """
  30. Select the transformation phases to use following args parameters.
  31. :return: <ImageTransform[]> list of image transformation
  32. """
  33. def shift_step(shift_starting=0, shift_ending=0):
  34. if not Conf.args['steps']:
  35. Conf.args['steps'] = (0, 5)
  36. Conf.args['steps'] = (
  37. Conf.args['steps'][0] + shift_starting,
  38. Conf.args['steps'][1] + shift_ending
  39. )
  40. def add_tail(phases, phase):
  41. phases = [phase] + phases
  42. if Conf.args['steps'] and Conf.args['steps'][0] != 0:
  43. shift_step(shift_starting=1)
  44. if Conf.args['steps'] and Conf.args['steps'][1] == len(phases) - 1:
  45. shift_step(shift_ending=1)
  46. return phases
  47. def add_head(phases, phase):
  48. phases = phases + [phase]
  49. if Conf.args['steps'] and Conf.args['steps'][1] == len(phases) - 1:
  50. shift_step(shift_ending=1)
  51. return phases
  52. phases = [DressToCorrect, CorrectToMask, MaskToMaskref,
  53. MaskrefToMaskdet, MaskdetToMaskfin, MaskfinToNude]
  54. if Conf.args['overlay']:
  55. phases = add_tail(phases, ImageToResized)
  56. phases = add_tail(phases, ImageToCrop)
  57. phases = add_head(phases, ImageToOverlay)
  58. elif Conf.args['auto_resize']:
  59. phases = add_tail(phases, ImageToResized)
  60. elif Conf.args['auto_resize_crop']:
  61. phases = add_tail(phases, ImageToResizedCrop)
  62. elif Conf.args['auto_rescale']:
  63. phases = add_tail(phases, ImageToRescale)
  64. elif os.path.isfile(Conf.args['input']):
  65. if not Conf.args['ignore_size']:
  66. check_shape(Conf.args['input'])
  67. else:
  68. Conf.log.warn('Image Size Requirements Unchecked.')
  69. if Conf.args['color_transfer']:
  70. phases = add_head(phases, ColorTransfer)
  71. return phases
  72. def select_processing():
  73. """
  74. Select the processing to use following args parameters.
  75. :return: <Process> a process to run
  76. """
  77. phases = select_phases()
  78. if os.path.isdir(Conf.args['input']):
  79. process = processing_image_folder(phases)
  80. elif Conf.args['n_runs'] != 1:
  81. process = multiple_image_processing(phases, Conf.args['n_runs'])
  82. else:
  83. process = simple_image_processing(phases)
  84. Conf.log.debug("Process to execute : {}".format(process))
  85. return process
  86. def simple_image_processing(phases):
  87. """
  88. Define a simple image process ready to run.
  89. :param phases: <ImageTransform[]> list of image transformation
  90. :return: <SimpleTransform> a image process run ready
  91. """
  92. return SimpleTransform(Conf.args['input'], phases, Conf.args['output'])
  93. def multiple_image_processing(phases, n_runs):
  94. """
  95. Define a multiple image process ready to run.
  96. :param phases: <ImageTransform[]> list of image transformation
  97. :param n_runs: number of times to process
  98. :return: <MultipleTransform> a multiple image process run ready
  99. """
  100. filename, extension = os.path.splitext(Conf.args['output'])
  101. return MultipleImageTransform(
  102. [Conf.args['input'] for _ in range(n_runs)],
  103. phases,
  104. ["{}{}{}".format(filename, i, extension) for i in range(n_runs)]
  105. )
  106. def processing_image_folder(phases):
  107. """
  108. Define a folder image process ready to run.
  109. :param phases: <ImageTransform[]> list of image transformation
  110. :return: <FolderImageTransform> a image process run ready
  111. """
  112. return FolderImageTransform(Conf.args['input'], phases, Conf.args['output'])
  113. if __name__ == "__main__":
  114. colorama.init()
  115. freeze_support()
  116. argv.run()