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.

inferrer.py 1.6KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566
  1. """Inference Body part functions."""
  2. import random
  3. from transform.opencv.bodypart import BodyPart
  4. def infer_nip(aur_list):
  5. """
  6. Infer nipples.
  7. :param aur_list: <BodyPart[]> aur list)
  8. :return: <BodyPart[]> nip list
  9. """
  10. nip_list = []
  11. for aur in aur_list:
  12. # Nip rules:
  13. # - circle (w == h)
  14. # - min dim: 5
  15. # - bigger if aur is bigger
  16. nip_dim = int(5 + aur.w * random.uniform(0.03, 0.09))
  17. # center:
  18. x = aur.x
  19. y = aur.y
  20. # Calculate Bounding Box:
  21. xmin = int(x - (nip_dim / 2))
  22. xmax = int(x + (nip_dim / 2))
  23. ymin = int(y - (nip_dim / 2))
  24. ymax = int(y + (nip_dim / 2))
  25. nip_list.append(BodyPart("nip", xmin, ymin, xmax, ymax, x, y, nip_dim, nip_dim))
  26. return nip_list
  27. def infer_hair(vag_list, enable):
  28. """
  29. Infer vaginal hair.
  30. :param vag_list: <BodyPart[]> vag list
  31. :param enable: <Boolean> Enable or disable hair generation
  32. :return: <BodyPart[]> hair list
  33. """
  34. hair_list = []
  35. if enable:
  36. for vag in vag_list:
  37. # Hair rules:
  38. hair_w = vag.w * random.uniform(0.4, 1.5)
  39. hair_h = vag.h * random.uniform(0.4, 1.5)
  40. # center:
  41. x = vag.x
  42. y = vag.y - (hair_h / 2) - (vag.h / 2)
  43. # Calculate Bounding Box:
  44. xmin = int(x - (hair_w / 2))
  45. xmax = int(x + (hair_w / 2))
  46. ymin = int(y - (hair_h / 2))
  47. ymax = int(y + (hair_h / 2))
  48. hair_list.append(BodyPart("hair", xmin, ymin, xmax, ymax, x, y, hair_w, hair_h))
  49. return hair_list