Browse Source

Merge pull request #45 from PommeDroid/more_refactor

More Refactor
tags/v1.2.10
deeppppp 2 years ago
parent
commit
b35a59983d
No account linked to committer's email address

+ 0
- 31
transform/opencv/__init__.py View File

@@ -13,34 +13,3 @@ class ImageTransformOpenCV(ImageTransform):
:param args: <dict> args parameter to run the image transformation (default use Conf.args)
"""
super().__init__(args=args, input_index=input_index)


class BodyPart:
"""Body part annotation."""

def __init__(self, name, xmin, ymin, xmax, ymax, x, y, w, h):
"""
Body Part constructor.

:param name: <string>
:param xmin: <int>
:param ymin: <int>
:param xmax: <int>
:param ymax: <int>
:param x: <int>
:param y: <int>
:param w: <int>
:param h: <int>
"""
self.name = name
# Bounding Box:
self.xmin = xmin
self.ymin = ymin
self.xmax = xmax
self.ymax = ymax
# Center:
self.x = x
self.y = y
# Dimensione:
self.w = w
self.h = h

+ 32
- 0
transform/opencv/bodypart/__init__.py View File

@@ -0,0 +1,32 @@
"""Body part Module."""


class BodyPart:
"""Body part annotation."""

def __init__(self, name, xmin, ymin, xmax, ymax, x, y, w, h):
"""
Body Part constructor.

:param name: <string>
:param xmin: <int>
:param ymin: <int>
:param xmax: <int>
:param ymax: <int>
:param x: <int>
:param y: <int>
:param w: <int>
:param h: <int>
"""
self.name = name
# Bounding Box:
self.xmin = xmin
self.ymin = ymin
self.xmax = xmax
self.ymax = ymax
# Center:
self.x = x
self.y = y
# Dimensione:
self.w = w
self.h = h

+ 198
- 0
transform/opencv/bodypart/extract.py View File

@@ -0,0 +1,198 @@
"""Extract Body part functions."""
import cv2
import numpy as np

from transform.opencv.bodypart import BodyPart
from transform.opencv.bodypart.inferrer import infer_nip, infer_hair
from transform.opencv.bodypart.resolver import detect_tit_aur_missing_problem, resolve_tit_aur_missing_problems


def extract_annotations(maskdet, enable_pubes):
"""
Extract Body part annotations.

:param maskdet: <RGB> maskdet image
:param enable_pubes: <Boolean> enable/disable pubic hair generation
:return: (<BodyPart []> bodypart_list) - for failure/error, return an empty list []
"""
# Find body part
tits_list = find_body_part(maskdet, "tit")
aur_list = find_body_part(maskdet, "aur")
vag_list = find_body_part(maskdet, "vag")
belly_list = find_body_part(maskdet, "belly")

# Filter out parts basing on dimension (area and aspect ratio):
aur_list = filter_dim_parts(aur_list, 100, 1000, 0.5, 3)
tits_list = filter_dim_parts(tits_list, 1000, 60000, 0.2, 3)
vag_list = filter_dim_parts(vag_list, 10, 1000, 0.2, 3)
belly_list = filter_dim_parts(belly_list, 10, 1000, 0.2, 3)

# Filter couple (if parts are > 2, choose only 2)
aur_list = filter_couple(aur_list)
tits_list = filter_couple(tits_list)

# Detect a missing problem (code of the problem)
missing_problem = detect_tit_aur_missing_problem(tits_list, aur_list)

# Check if problem is SOLVEABLE:
if missing_problem in [3, 6, 7, 8]:
resolve_tit_aur_missing_problems(tits_list, aur_list, missing_problem)

# Infer the nips:
nip_list = infer_nip(aur_list)

# Infer the hair:
hair_list = infer_hair(vag_list, enable_pubes)

# Return a combined list:
return tits_list + aur_list + nip_list + vag_list + hair_list + belly_list


def find_body_part(image, part_name):
"""
Find body part.

:param image: <RGB> image
:param part_name: <string> part_name
:return: <BodyPart[]>list
"""
def calculate_bounding_box(h, w, x, y):
"""Calculate Bounding Box."""
xmin = int(x - (w / 2))
xmax = int(x + (w / 2))
ymin = int(y - (h / 2))
ymax = int(y + (h / 2))
return xmax, xmin, ymax, ymin

def detect_direction(a_max, a_min, angle):
"""Detect direction."""
if angle == 0:
h = a_max
w = a_min
else:
h = a_min
w = a_max
return h, w

def normalize_belly_vag(h, part_name, w):
"""Normalize the belly and vag size."""
if part_name in ("belly", "vag"):
if w < 15:
w *= 2
if h < 15:
h *= 2
return h, w

def get_correct_filter_color(image, part_name):
"""Get the correct color filter."""
if part_name == "tit":
# Use combined color filter
f1 = np.asarray([0, 0, 0]) # tit color filter
f2 = np.asarray([10, 10, 10])
f3 = np.asarray([0, 0, 250]) # aur color filter
f4 = np.asarray([0, 0, 255])
color_mask1 = cv2.inRange(image, f1, f2)
color_mask2 = cv2.inRange(image, f3, f4)
color_mask = cv2.bitwise_or(color_mask1, color_mask2) # combine

elif part_name == "aur":
f1 = np.asarray([0, 0, 250]) # aur color filter
f2 = np.asarray([0, 0, 255])
color_mask = cv2.inRange(image, f1, f2)

elif part_name == "vag":
f1 = np.asarray([250, 0, 0]) # vag filter
f2 = np.asarray([255, 0, 0])
color_mask = cv2.inRange(image, f1, f2)

elif part_name == "belly":
f1 = np.asarray([250, 0, 250]) # belly filter
f2 = np.asarray([255, 0, 255])
color_mask = cv2.inRange(image, f1, f2)

return color_mask

bodypart_list = [] # empty BodyPart list

color_mask = get_correct_filter_color(image, part_name)

# find contours:
contours, _ = cv2.findContours(color_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

# for every contour:
for cnt in contours:

if len(cnt) > 5: # at least 5 points to fit ellipse

# (x, y), (MA, ma), angle = cv2.fitEllipse(cnt)
ellipse = cv2.fitEllipse(cnt)

# Fit Result:
x = ellipse[0][0] # center x
y = ellipse[0][1] # center y
angle = ellipse[2] # angle
a_min = ellipse[1][0] # asse minore
a_max = ellipse[1][1] # asse maggiore

h, w = detect_direction(a_max, a_min, angle)

h, w = normalize_belly_vag(h, part_name, w)

xmax, xmin, ymax, ymin = calculate_bounding_box(h, w, x, y)

bodypart_list.append(BodyPart(part_name, xmin, ymin, xmax, ymax, x, y, w, h))

return bodypart_list


def filter_dim_parts(bp_list, min_area, max_area, min_ar, max_ar):
"""
Filter a body part list with area and aspect ration.

:param bp_list: BodyPart[]>list
:param min_area: <num> minimum area of part
:param max_area: <num> max area
:param min_ar: <num> min aspect ratio
:param max_ar: <num> max aspect ratio
:return: <BodyPart[]>list
"""
b_filt = []

for obj in bp_list:
if min_area < obj.w * obj.h < max_area and min_ar < obj.w / obj.h < max_ar:
b_filt.append(obj)

return b_filt


def filter_couple(bp_list):
"""
Filer couple in body part list.

:param bp_list: <BodyPart[]>list
:return: <BodyPart[]>list
"""
# Remove exceed parts
if len(bp_list) > 2:

# trovare coppia (a,b) che minimizza bp_list[a].y-bp_list[b].y
min_a = 0
min_b = 1
min_diff = abs(bp_list[min_a].y - bp_list[min_b].y)

for a, _ in enumerate(bp_list):
for b, _ in enumerate(bp_list):
# TODO: avoid repetition (1,0) (0,1)
if a != b:
diff = abs(bp_list[a].y - bp_list[b].y)
if diff < min_diff:
min_diff = diff
min_a = a
min_b = b

b_filt = [bp_list[min_a], bp_list[min_b]]

return b_filt
else:
# No change
return bp_list

+ 66
- 0
transform/opencv/bodypart/inferrer.py View File

@@ -0,0 +1,66 @@
"""Inference Body part functions."""
import random

from transform.opencv.bodypart import BodyPart


def infer_nip(aur_list):
"""
Infer nipples.

:param aur_list: <BodyPart[]> aur list)
:return: <BodyPart[]> nip list
"""
nip_list = []

for aur in aur_list:
# Nip rules:
# - circle (w == h)
# - min dim: 5
# - bigger if aur is bigger
nip_dim = int(5 + aur.w * random.uniform(0.03, 0.09))

# center:
x = aur.x
y = aur.y

# Calculate Bounding Box:
xmin = int(x - (nip_dim / 2))
xmax = int(x + (nip_dim / 2))
ymin = int(y - (nip_dim / 2))
ymax = int(y + (nip_dim / 2))

nip_list.append(BodyPart("nip", xmin, ymin, xmax, ymax, x, y, nip_dim, nip_dim))

return nip_list


def infer_hair(vag_list, enable):
"""
Infer vaginal hair.

:param vag_list: <BodyPart[]> vag list
:param enable: <Boolean> Enable or disable hair generation
:return: <BodyPart[]> hair list
"""
hair_list = []

if enable:
for vag in vag_list:
# Hair rules:
hair_w = vag.w * random.uniform(0.4, 1.5)
hair_h = vag.h * random.uniform(0.4, 1.5)

# center:
x = vag.x
y = vag.y - (hair_h / 2) - (vag.h / 2)

# Calculate Bounding Box:
xmin = int(x - (hair_w / 2))
xmax = int(x + (hair_w / 2))
ymin = int(y - (hair_h / 2))
ymax = int(y + (hair_h / 2))

hair_list.append(BodyPart("hair", xmin, ymin, xmax, ymax, x, y, hair_w, hair_h))

return hair_list

+ 176
- 0
transform/opencv/bodypart/resolver.py View File

@@ -0,0 +1,176 @@
"""Inference Body problems resolver."""
import random

from transform.opencv.bodypart import BodyPart


def detect_tit_aur_missing_problem(tits_list, aur_list):
"""
Detect tits aur missing problem.

(<int> problem code)
# TIT | AUR | code | SOLVE? |
# 0 | 0 | 1 | NO |
# 0 | 1 | 2 | NO |
# 0 | 2 | 3 | YES |
# 1 | 0 | 4 | NO |
# 1 | 1 | 5 | NO |
# 1 | 2 | 6 | YES |
# 2 | 0 | 7 | YES |
# 2 | 1 | 8 | YES |

:param tits_list: <BodyPart[]> tits list
:param aur_list: <BodyPart[]> aur list
:return: <int> problem code
"""
return {
(0, 0): 1,
(0, 1): 2,
(0, 2): 3,
(1, 0): 4,
(1, 1): 5,
(1, 2): 6,
(2, 0): 7,
(2, 1): 8,
}.get((len(tits_list), len(aur_list)), -1)


def resolve_tit_aur_missing_problems(tits_list, aur_list, problem_code):
"""
Resolve tits missing aur problem.

:param tits_list: <BodyPart[]> tits list
:param aur_list: <BodyPart[]> aur list
:param problem_code: <int> problem code
:return: None
"""
def resolve_problem_3():
random_tit_factor = random.randint(2, 5) # TOTEST

# Add the first tit:
new_w = aur_list[0].w * random_tit_factor # TOTEST
new_x = aur_list[0].x
new_y = aur_list[0].y

xmin = int(new_x - (new_w / 2))
xmax = int(new_x + (new_w / 2))
ymin = int(new_y - (new_w / 2))
ymax = int(new_y + (new_w / 2))

tits_list.append(BodyPart("tit", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w))

# Add the second tit:
new_w = aur_list[1].w * random_tit_factor # TOTEST
new_x = aur_list[1].x
new_y = aur_list[1].y

xmin = int(new_x - (new_w / 2))
xmax = int(new_x + (new_w / 2))
ymin = int(new_y - (new_w / 2))
ymax = int(new_y + (new_w / 2))

tits_list.append(BodyPart("tit", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w))

def resolve_problem_6():
# Find witch aur is full:
d1 = abs(tits_list[0].x - aur_list[0].x)
d2 = abs(tits_list[0].x - aur_list[1].x)

if d1 > d2:
# aur[0] is empty
new_x = aur_list[0].x
new_y = aur_list[0].y
else:
# aur[1] is empty
new_x = aur_list[1].x
new_y = aur_list[1].y

# Calculate Bounding Box:
xmin = int(new_x - (tits_list[0].w / 2))
xmax = int(new_x + (tits_list[0].w / 2))
ymin = int(new_y - (tits_list[0].w / 2))
ymax = int(new_y + (tits_list[0].w / 2))

tits_list.append(BodyPart("tit", xmin, ymin, xmax, ymax, new_x, new_y, tits_list[0].w, tits_list[0].w))

def resolve_problem_7():
# Add the first aur:
new_w = tits_list[0].w * random.uniform(0.03, 0.1) # TOTEST
new_x = tits_list[0].x
new_y = tits_list[0].y

xmin = int(new_x - (new_w / 2))
xmax = int(new_x + (new_w / 2))
ymin = int(new_y - (new_w / 2))
ymax = int(new_y + (new_w / 2))

aur_list.append(BodyPart("aur", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w))

# Add the second aur:
new_w = tits_list[1].w * random.uniform(0.03, 0.1) # TOTEST
new_x = tits_list[1].x
new_y = tits_list[1].y

xmin = int(new_x - (new_w / 2))
xmax = int(new_x + (new_w / 2))
ymin = int(new_y - (new_w / 2))
ymax = int(new_y + (new_w / 2))

aur_list.append(BodyPart("aur", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w))

def resolve_problem_8():
# Find witch tit is full:
d1 = abs(aur_list[0].x - tits_list[0].x)
d2 = abs(aur_list[0].x - tits_list[1].x)

if d1 > d2:
# tit[0] is empty
new_x = tits_list[0].x
new_y = tits_list[0].y
else:
# tit[1] is empty
new_x = tits_list[1].x
new_y = tits_list[1].y

# Calculate Bounding Box:
xmin = int(new_x - (aur_list[0].w / 2))
xmax = int(new_x + (aur_list[0].w / 2))
ymin = int(new_y - (aur_list[0].w / 2))
ymax = int(new_y + (aur_list[0].w / 2))
aur_list.append(BodyPart("aur", xmin, ymin, xmax, ymax, new_x, new_y, aur_list[0].w, aur_list[0].w))

{
3: resolve_problem_3,
6: resolve_problem_6,
7: resolve_problem_7,
8: resolve_problem_8,
}.get(problem_code, lambda _: _)()


def detect_tit_aur_position_problem(tits_list, aur_list):
"""
Detect tits position problem.

:param tits_list: <BodyPart[]> tits list
:param aur_list: <BodyPart[]> aur list
:return: <Boolean>
"""
def detect_tits_too_narrow_horizontally():
diff_tits_x = abs(tits_list[0].x - tits_list[1].x)
return diff_tits_x < 40

def detect_tits_too_narrow_vertically():
diff_tits_y = abs(tits_list[0].y - tits_list[1].y)
return diff_tits_y > 120

def detect_tits_too_equal_or_different_width():
diff_tits_w = abs(tits_list[0].w - tits_list[1].w)
return (diff_tits_w < 0.1) or (diff_tits_w > 60)

def detect_tits_body_position_is_too_low():
# Calculate the ratio between y and aurs distance
rapp = aur_list[0].y / (abs(aur_list[0].x - aur_list[1].x))
return aur_list[0].y > 350 and rapp > 2.8

return (detect_tits_too_narrow_horizontally() or detect_tits_too_narrow_vertically() or
detect_tits_too_equal_or_different_width() or detect_tits_body_position_is_too_low)

+ 48
- 485
transform/opencv/mask.py View File

@@ -1,18 +1,18 @@
"""OpenCV Mask Transforms."""
import random

import cv2
import numpy as np

from transform.opencv import ImageTransformOpenCV, BodyPart
from transform.opencv import ImageTransformOpenCV
from transform.opencv.bodypart.extract import extract_annotations


class MaskToMaskref(ImageTransformOpenCV):
"""Mask & Correct -> MaskRef [OPENCV]."""
class MaskImageTransformOpenCV(ImageTransformOpenCV):
"""Mask Image Transform OpenCV."""

def __init__(self, input_index=(-2, -1), args=None):
"""
Mask To Maskref constructor.
Mask Image Transform OpenCV constructor.

:param input_index: <tuple> index where to take the inputs (default is (-2,-1)
for the two previous transformation)
@@ -20,6 +20,10 @@ class MaskToMaskref(ImageTransformOpenCV):
"""
super().__init__(args=args, input_index=input_index)


class MaskToMaskref(MaskImageTransformOpenCV):
"""Mask & Correct -> MaskRef [OPENCV]."""

def _execute(self, *args):
"""
Create mask ref.
@@ -54,7 +58,7 @@ class MaskToMaskref(ImageTransformOpenCV):
return cv2.add(res1, res2)


class MaskdetToMaskfin(ImageTransformOpenCV):
class MaskdetToMaskfin(MaskImageTransformOpenCV):
"""Maskdet -> Maskfin [OPENCV]."""

def __init__(self, input_index=(-2, -1), args=None,):
@@ -91,18 +95,7 @@ class MaskdetToMaskfin(ImageTransformOpenCV):
def to_int(a, b):
return int(round(a * float(b)))

enable_pubes = (self.__hair_size > 0)

# Create a total green image, in which draw details ellipses
details = np.zeros((512, 512, 3), np.uint8)
details[:, :, :] = (0, 255, 0) # (B, G, R)

# Extract body part features:
bodypart_list = self.extract_annotations(args[1], enable_pubes)

# Check if the list is not empty:
if bodypart_list:

def draw_bodypart_details(bodypart_list, details, to_int):
# Draw body part in details image:
for obj in bodypart_list:

@@ -129,23 +122,43 @@ class MaskdetToMaskfin(ImageTransformOpenCV):
hairmax = to_int(self.__hair_size, a_max)
hairmin = to_int(self.__hair_size, a_min)

# Draw ellipse
if obj.name == "tit":
cv2.ellipse(details, (x, y), (titmax, titmin), angle, 0, 360, (0, 205, 0), -1) # (0,0,0,50)
elif obj.name == "aur":
cv2.ellipse(details, (x, y), (aurmax, aurmin), angle, 0, 360, (0, 0, 255), -1) # red
elif obj.name == "nip":
cv2.ellipse(details, (x, y), (nipmax, nipmin), angle, 0, 360, (255, 255, 255), -1) # white
elif obj.name == "belly":
cv2.ellipse(details, (x, y), (a_max, a_min), angle, 0, 360, (255, 0, 255), -1) # purple
elif obj.name == "vag":
cv2.ellipse(details, (x, y), (vagmax, vagmin), angle, 0, 360, (255, 0, 0), -1) # blue
elif obj.name == "hair":
xmin = x - hairmax
ymin = y - hairmin
xmax = x + hairmax
ymax = y + hairmax
cv2.rectangle(details, (xmin, ymin), (xmax, ymax), (100, 100, 100), -1)
draw_ellipse(a_max, a_min, angle, aurmax, aurmin, details, hairmax, hairmin, nipmax, nipmin, obj,
titmax, titmin, vagmax, vagmin, x, y)

def draw_ellipse(a_max, a_min, angle, aurmax, aurmin, details, hairmax, hairmin, nipmax, nipmin, obj,
titmax,
titmin, vagmax, vagmin, x, y):
# Draw ellipse
if obj.name == "tit":
cv2.ellipse(details, (x, y), (titmax, titmin), angle, 0, 360, (0, 205, 0), -1) # (0,0,0,50)
elif obj.name == "aur":
cv2.ellipse(details, (x, y), (aurmax, aurmin), angle, 0, 360, (0, 0, 255), -1) # red
elif obj.name == "nip":
cv2.ellipse(details, (x, y), (nipmax, nipmin), angle, 0, 360, (255, 255, 255), -1) # white
elif obj.name == "belly":
cv2.ellipse(details, (x, y), (a_max, a_min), angle, 0, 360, (255, 0, 255), -1) # purple
elif obj.name == "vag":
cv2.ellipse(details, (x, y), (vagmax, vagmin), angle, 0, 360, (255, 0, 0), -1) # blue
elif obj.name == "hair":
xmin = x - hairmax
ymin = y - hairmin
xmax = x + hairmax
ymax = y + hairmax
cv2.rectangle(details, (xmin, ymin), (xmax, ymax), (100, 100, 100), -1)

enable_pubes = (self.__hair_size > 0)

# Create a total green image, in which draw details ellipses
details = np.zeros((512, 512, 3), np.uint8)
details[:, :, :] = (0, 255, 0) # (B, G, R)

# Extract body part features:
bodypart_list = extract_annotations(args[1], enable_pubes)

# Check if the list is not empty:
if bodypart_list:

draw_bodypart_details(bodypart_list, details, to_int)

# Define the green color filter
f1 = np.asarray([0, 250, 0]) # green color filter
@@ -164,453 +177,3 @@ class MaskdetToMaskfin(ImageTransformOpenCV):
# Compone:
maskfin = cv2.add(res1, res2)
return maskfin

def extract_annotations(self, maskdet, enable_pubes):
"""
Extract annotations.

:param maskdet: <RGB> maskdet image
:param enable_pubes: <Boolean> enable/disable pubic hair generation
:return: (<BodyPart []> bodypart_list) - for failure/error, return an empty list []
"""
# Find body part
tits_list = self.find_body_part(maskdet, "tit")
aur_list = self.find_body_part(maskdet, "aur")
vag_list = self.find_body_part(maskdet, "vag")
belly_list = self.find_body_part(maskdet, "belly")

# Filter out parts basing on dimension (area and aspect ratio):
aur_list = self.filter_dim_parts(aur_list, 100, 1000, 0.5, 3)
tits_list = self.filter_dim_parts(tits_list, 1000, 60000, 0.2, 3)
vag_list = self.filter_dim_parts(vag_list, 10, 1000, 0.2, 3)
belly_list = self.filter_dim_parts(belly_list, 10, 1000, 0.2, 3)

# Filter couple (if parts are > 2, choose only 2)
aur_list = self.filter_couple(aur_list)
tits_list = self.filter_couple(tits_list)

# Detect a missing problem (code of the problem)
missing_problem = self.detect_tit_aur_missing_problem(tits_list, aur_list)

# Check if problem is SOLVEABLE:
if missing_problem in [3, 6, 7, 8]:
self.resolve_tit_aur_missing_problems(tits_list, aur_list, missing_problem)

# Infer the nips:
nip_list = self.infer_nip(aur_list)

# Infer the hair:
hair_list = self.infer_hair(vag_list, enable_pubes)

# Return a combined list:
return tits_list + aur_list + nip_list + vag_list + hair_list + belly_list

@staticmethod
def find_body_part(image, part_name):
"""
Find body part.

:param image: <RGB> image
:param part_name: <string> part_name
:return: <BodyPart[]>list
"""
bodypart_list = [] # empty BodyPart list

# Get the correct color filter:
if part_name == "tit":
# Use combined color filter
f1 = np.asarray([0, 0, 0]) # tit color filter
f2 = np.asarray([10, 10, 10])
f3 = np.asarray([0, 0, 250]) # aur color filter
f4 = np.asarray([0, 0, 255])
color_mask1 = cv2.inRange(image, f1, f2)
color_mask2 = cv2.inRange(image, f3, f4)
color_mask = cv2.bitwise_or(color_mask1, color_mask2) # combine

elif part_name == "aur":
f1 = np.asarray([0, 0, 250]) # aur color filter
f2 = np.asarray([0, 0, 255])
color_mask = cv2.inRange(image, f1, f2)

elif part_name == "vag":
f1 = np.asarray([250, 0, 0]) # vag filter
f2 = np.asarray([255, 0, 0])
color_mask = cv2.inRange(image, f1, f2)

elif part_name == "belly":
f1 = np.asarray([250, 0, 250]) # belly filter
f2 = np.asarray([255, 0, 255])
color_mask = cv2.inRange(image, f1, f2)

# find contours:
contours, _ = cv2.findContours(color_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

# for every contour:
for cnt in contours:

if len(cnt) > 5: # at least 5 points to fit ellipse

# (x, y), (MA, ma), angle = cv2.fitEllipse(cnt)
ellipse = cv2.fitEllipse(cnt)

# Fit Result:
x = ellipse[0][0] # center x
y = ellipse[0][1] # center y
angle = ellipse[2] # angle
a_min = ellipse[1][0] # asse minore
a_max = ellipse[1][1] # asse maggiore

# Detect direction:
if angle == 0:
h = a_max
w = a_min
else:
h = a_min
w = a_max

# Normalize the belly size:
if part_name == "belly":
if w < 15:
w *= 2
if h < 15:
h *= 2

# Normalize the vag size:
if part_name == "vag":
if w < 15:
w *= 2
if h < 15:
h *= 2

# Calculate Bounding Box:
xmin = int(x - (w / 2))
xmax = int(x + (w / 2))
ymin = int(y - (h / 2))
ymax = int(y + (h / 2))

bodypart_list.append(BodyPart(part_name, xmin, ymin, xmax, ymax, x, y, w, h))

return bodypart_list

@staticmethod
def filter_dim_parts(bp_list, min_area, max_area, min_ar, max_ar):
"""
Filter a body part list with area and aspect ration.

:param bp_list: BodyPart[]>list
:param min_area: <num> minimum area of part
:param max_area: <num> max area
:param min_ar: <num> min aspect ratio
:param max_ar: <num> max aspect ratio
:return: <BodyPart[]>list
"""
b_filt = []

for obj in bp_list:

a = obj.w * obj.h # Object AREA

if ((a > min_area) and (a < max_area)):

ar = obj.w / obj.h # Object ASPECT RATIO

if ((ar > min_ar) and (ar < max_ar)):
b_filt.append(obj)

return b_filt

@staticmethod
def filter_couple(bp_list):
"""
Filer couple in body part list.

:param bp_list: <BodyPart[]>list
:return: <BodyPart[]>list
"""
# Remove exceed parts
if len(bp_list) > 2:

# trovare coppia (a,b) che minimizza bp_list[a].y-bp_list[b].y
min_a = 0
min_b = 1
min_diff = abs(bp_list[min_a].y - bp_list[min_b].y)

for a, _ in enumerate(bp_list):
for b, _ in enumerate(bp_list):
# TODO: avoid repetition (1,0) (0,1)
if a != b:
diff = abs(bp_list[a].y - bp_list[b].y)
if diff < min_diff:
min_diff = diff
min_a = a
min_b = b
b_filt = []

b_filt.append(bp_list[min_a])
b_filt.append(bp_list[min_b])

return b_filt
else:
# No change
return bp_list

@staticmethod
def detect_tit_aur_missing_problem(tits_list, aur_list):
"""
Detect tits aur missing problem.

(<int> problem code)
# TIT | AUR | code | SOLVE? |
# 0 | 0 | 1 | NO |
# 0 | 1 | 2 | NO |
# 0 | 2 | 3 | YES |
# 1 | 0 | 4 | NO |
# 1 | 1 | 5 | NO |
# 1 | 2 | 6 | YES |
# 2 | 0 | 7 | YES |
# 2 | 1 | 8 | YES |

:param tits_list: <BodyPart[]> tits list
:param aur_list: <BodyPart[]> aur list
:return: <int> problem code
"""
t_len = len(tits_list)
a_len = len(aur_list)

if t_len == 0:
if a_len == 0:
return 1
elif a_len == 1:
return 2
elif a_len == 2:
return 3
else:
return -1
elif t_len == 1:
if a_len == 0:
return 4
elif a_len == 1:
return 5
elif a_len == 2:
return 6
else:
return -1
elif t_len == 2:
if a_len == 0:
return 7
elif a_len == 1:
return 8
else:
return -1
else:
return -1

@staticmethod
def resolve_tit_aur_missing_problems(tits_list, aur_list, problem_code):
"""
Resolve tits missing aur problem.

:param tits_list: <BodyPart[]> tits list
:param aur_list: <BodyPart[]> aur list
:param problem_code: <int> problem code
:return: None
"""
if problem_code == 3:

random_tit_factor = random.randint(2, 5) # TOTEST

# Add the first tit:
new_w = aur_list[0].w * random_tit_factor # TOTEST
new_x = aur_list[0].x
new_y = aur_list[0].y

xmin = int(new_x - (new_w / 2))
xmax = int(new_x + (new_w / 2))
ymin = int(new_y - (new_w / 2))
ymax = int(new_y + (new_w / 2))

tits_list.append(BodyPart("tit", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w))

# Add the second tit:
new_w = aur_list[1].w * random_tit_factor # TOTEST
new_x = aur_list[1].x
new_y = aur_list[1].y

xmin = int(new_x - (new_w / 2))
xmax = int(new_x + (new_w / 2))
ymin = int(new_y - (new_w / 2))
ymax = int(new_y + (new_w / 2))

tits_list.append(BodyPart("tit", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w))

elif problem_code == 6:

# Find wich aur is full:
d1 = abs(tits_list[0].x - aur_list[0].x)
d2 = abs(tits_list[0].x - aur_list[1].x)

if d1 > d2:
# aur[0] is empty
new_x = aur_list[0].x
new_y = aur_list[0].y
else:
# aur[1] is empty
new_x = aur_list[1].x
new_y = aur_list[1].y

# Calculate Bounding Box:
xmin = int(new_x - (tits_list[0].w / 2))
xmax = int(new_x + (tits_list[0].w / 2))
ymin = int(new_y - (tits_list[0].w / 2))
ymax = int(new_y + (tits_list[0].w / 2))

tits_list.append(BodyPart("tit", xmin, ymin, xmax, ymax, new_x, new_y, tits_list[0].w, tits_list[0].w))

elif problem_code == 7:

# Add the first aur:
new_w = tits_list[0].w * random.uniform(0.03, 0.1) # TOTEST
new_x = tits_list[0].x
new_y = tits_list[0].y

xmin = int(new_x - (new_w / 2))
xmax = int(new_x + (new_w / 2))
ymin = int(new_y - (new_w / 2))
ymax = int(new_y + (new_w / 2))

aur_list.append(BodyPart("aur", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w))

# Add the second aur:
new_w = tits_list[1].w * random.uniform(0.03, 0.1) # TOTEST
new_x = tits_list[1].x
new_y = tits_list[1].y

xmin = int(new_x - (new_w / 2))
xmax = int(new_x + (new_w / 2))
ymin = int(new_y - (new_w / 2))
ymax = int(new_y + (new_w / 2))

aur_list.append(BodyPart("aur", xmin, ymin, xmax, ymax, new_x, new_y, new_w, new_w))

elif problem_code == 8:

# Find wich tit is full:
d1 = abs(aur_list[0].x - tits_list[0].x)
d2 = abs(aur_list[0].x - tits_list[1].x)

if d1 > d2:
# tit[0] is empty
new_x = tits_list[0].x
new_y = tits_list[0].y
else:
# tit[1] is empty
new_x = tits_list[1].x
new_y = tits_list[1].y

# Calculate Bounding Box:
xmin = int(new_x - (aur_list[0].w / 2))
xmax = int(new_x + (aur_list[0].w / 2))
ymin = int(new_y - (aur_list[0].w / 2))
ymax = int(new_y + (aur_list[0].w / 2))
aur_list.append(BodyPart("aur", xmin, ymin, xmax, ymax, new_x, new_y, aur_list[0].w, aur_list[0].w))

@staticmethod
def detect_tit_aur_position_problem(tits_list, aur_list):
"""
Detect tits position problem.

:param tits_list: <BodyPart[]> tits list
:param aur_list: <BodyPart[]> aur list
:return: <Boolean>
"""
diff_tits_x = abs(tits_list[0].x - tits_list[1].x)
if diff_tits_x < 40:
print("diffTitsX")
# Tits too narrow (orizontally)
return True

diff_tits_y = abs(tits_list[0].y - tits_list[1].y)
if diff_tits_y > 120:
# Tits too distanced (vertically)
print("diffTitsY")
return True

diff_tits_w = abs(tits_list[0].w - tits_list[1].w)
if (diff_tits_w < 0.1) or (diff_tits_w > 60):
print("diffTitsW")
# Tits too equals, or too different (width)
return True

# Check if body position is too low (face not covered by watermark)
if aur_list[0].y > 350: # tits too low
# Calculate the ratio between y and aurs distance
rapp = aur_list[0].y / (abs(aur_list[0].x - aur_list[1].x))
if rapp > 2.8:
print("aurDown")
return True

return False

@staticmethod
def infer_nip(aur_list):
"""
Infer nipples.

:param aur_list: <BodyPart[]> aur list)
:return: <BodyPart[]> nip list
"""
nip_list = []

for aur in aur_list:
# Nip rules:
# - circle (w == h)
# - min dim: 5
# - bigger if aur is bigger
nip_dim = int(5 + aur.w * random.uniform(0.03, 0.09))

# center:
x = aur.x
y = aur.y

# Calculate Bounding Box:
xmin = int(x - (nip_dim / 2))
xmax = int(x + (nip_dim / 2))
ymin = int(y - (nip_dim / 2))
ymax = int(y + (nip_dim / 2))

nip_list.append(BodyPart("nip", xmin, ymin, xmax, ymax, x, y, nip_dim, nip_dim))

return nip_list

@staticmethod
def infer_hair(vag_list, enable):
"""
Infer vaginal hair.

:param vag_list: <BodyPart[]> vag list
:param enable: <Boolean> Enable or disable hair generation
:return: <BodyPart[]> hair list
"""
hair_list = []

# 70% of chanche to add hair
if enable:

for vag in vag_list:
# Hair rules:
hair_w = vag.w * random.uniform(0.4, 1.5)
hair_h = vag.h * random.uniform(0.4, 1.5)

# center:
x = vag.x
y = vag.y - (hair_h / 2) - (vag.h / 2)

# Calculate Bounding Box:
xmin = int(x - (hair_w / 2))
xmax = int(x + (hair_w / 2))
ymin = int(y - (hair_h / 2))
ymax = int(y + (hair_h / 2))

hair_list.append(BodyPart("hair", xmin, ymin, xmax, ymax, x, y, hair_w, hair_h))

return hair_list

Loading…
Cancel
Save