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.3KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293
  1. import os
  2. from config import Config as Conf
  3. from transform.gan.mask import CorrectToMask, MaskrefToMaskdet, MaskfinToNude
  4. from transform.opencv.correct import DressToCorrect, ColorTransfer
  5. from transform.opencv.mask import MaskToMaskref, MaskdetToMaskfin
  6. from transform.opencv.resize import ImageToResized, ImageToCrop, ImageToOverlay, ImageToResizedCrop, ImageToRescale
  7. from utils import check_shape
  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):
  42. if not args['ignore_size']:
  43. check_shape(args['input'])
  44. else:
  45. Conf.log.warn('Image Size Requirements Unchecked.')
  46. def scale_mod(args, p):
  47. for mod in (overlay, auto_resize, auto_resize_crop, auto_rescale):
  48. if args.get(mod.__name__):
  49. return mod(args, p)
  50. if os.path.isfile(Conf.args["input"]):
  51. is_file(args)
  52. return p
  53. def select_phases(args):
  54. """
  55. Select the transformation phases to use following args parameters.
  56. :return: <ImageTransform[]> list of image transformation
  57. """
  58. phases = [DressToCorrect, CorrectToMask, MaskToMaskref,
  59. MaskrefToMaskdet, MaskdetToMaskfin, MaskfinToNude]
  60. phases = scale_mod(args, phases)
  61. if args['color_transfer']:
  62. phases = add_head(args, phases, ColorTransfer)
  63. return phases