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.

utils.py 2.5KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101
  1. import os
  2. from config import Config as Conf
  3. from utils import check_shape
  4. from transform.gan.mask import CorrectToMask, MaskrefToMaskdet, MaskfinToNude
  5. from transform.opencv.correct import DressToCorrect, ColorTransfer
  6. from transform.opencv.mask import MaskToMaskref, MaskdetToMaskfin
  7. from transform.opencv.resize import ImageToResized, ImageToCrop, ImageToOverlay, ImageToResizedCrop, ImageToRescale
  8. def shift_step(args, shift_start_add=0, shift_end_add=0):
  9. if not args['steps']:
  10. args['steps'] = (0, 5)
  11. args['steps'] = (
  12. args['steps'][0] + shift_start_add,
  13. args['steps'][1] + shift_end_add
  14. )
  15. def shift_starting(args):
  16. if args.get('steps') and args['steps'][0] != 0:
  17. shift_step(args, shift_start_add=1)
  18. def shift_ending(args, p):
  19. if args.get('steps') and args['steps'][1] == len(p) - 1:
  20. shift_step(args, shift_end_add=1)
  21. def add_tail(args, p, add):
  22. p = [add] + p
  23. shift_starting(args)
  24. shift_ending(args, p)
  25. return p
  26. def add_head(args, p, add):
  27. p = p + [add]
  28. shift_ending(args, p)
  29. return p
  30. def overlay(args, p):
  31. p = add_tail(args, p, ImageToResized)
  32. p = add_tail(args, p, ImageToCrop)
  33. p = add_head(args, p, ImageToOverlay)
  34. return p
  35. def auto_resize(args, p):
  36. return add_tail(args, p, ImageToResized)
  37. def auto_resize_crop(args, p):
  38. return add_tail(args, p, ImageToResizedCrop)
  39. def auto_rescale(args, p):
  40. return add_tail(args, p, ImageToRescale)
  41. def is_file(args, path):
  42. if not os.path.isfile(path):
  43. return False
  44. for mod in (overlay, auto_resize, auto_resize_crop, auto_rescale):
  45. if args.get(mod.__name__):
  46. return True
  47. if not args['ignore_size']:
  48. check_shape(path)
  49. else:
  50. Conf.log.warn('Image Size Requirements Unchecked.')
  51. return True
  52. def scale_mod(args, p):
  53. for mod in (overlay, auto_resize, auto_resize_crop, auto_rescale):
  54. if args.get(mod.__name__):
  55. return mod(args, p)
  56. if os.path.isfile(Conf.args["input"]):
  57. is_file(args, Conf.args["input"])
  58. return p
  59. def select_phases(args):
  60. """
  61. Select the transformation phases to use following args parameters.
  62. :return: <ImageTransform[]> list of image transformation
  63. """
  64. phases = [DressToCorrect, CorrectToMask, MaskToMaskref,
  65. MaskrefToMaskdet, MaskdetToMaskfin, MaskfinToNude]
  66. phases = scale_mod(args, phases)
  67. if args['color_transfer']:
  68. phases = add_head(args, phases, ColorTransfer)
  69. return phases