added node to detect trouser images

This commit is contained in:
Ubuntu
2024-08-22 12:55:16 +00:00
parent a6cfa6483e
commit ff2cd75e6f
2 changed files with 40 additions and 4 deletions
+6 -4
View File
@@ -15,7 +15,7 @@ from scaled_paste import main_scaled_paste
from scaled_paste import main_scaled_paste_2
from simple_bg_swap import (simple_bg_swap, get_threshold_for_bg_swap, RGB_2_LAB, LAB_2_RGB, get_mean_and_standard_deviation, renormalize_array)
from distribution_reshape import (simple_rescale_histogram, get_histogram_limits)
from utility_nodes import TRI3D_clean_mask, TRI3D_extract_pose_part, TRI3D_position_pose_part, TRI3D_fill_mask
from utility_nodes import TRI3D_clean_mask, TRI3D_extract_pose_part, TRI3D_position_pose_part, TRI3D_fill_mask, TRI3D_is_only_trouser
from .AEMatter import (load_AEMatter_Model, run_AEMatter_inference)
from .light_layer import main_light_layer
@@ -3730,11 +3730,12 @@ NODE_CLASS_MAPPINGS = {
'tri3d-clean_mask': TRI3D_clean_mask,
"tri3d-extract_pose_part": TRI3D_extract_pose_part,
"tri3d_position_pose_part":TRI3D_position_pose_part,
"tri3d_fill_mask": TRI3D_fill_mask
"tri3d_fill_mask": TRI3D_fill_mask,
"tri3d_is_only_trouser": TRI3D_is_only_trouser
}
VERSION = "4.7.1"
VERSION = "4.7.2"
# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
"tri3d-photoroom-bgremove-api": "Photoroom BG Remove" + " v" + VERSION,
@@ -3788,5 +3789,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
'tri3d-clean_mask': "Clear small patches" + " v" + VERSION,
"tri3d-extract_pose_part": "Extract pose part" + " v" + VERSION,
"tri3d_position_pose_part": "Position pose part" + " v" + VERSION,
"tri3d_fill_mask": "Fill mask" + " v" + VERSION
"tri3d_fill_mask": "Fill mask" + " v" + VERSION,
"tri3d_is_only_trouser": "Is only trouser" + " v" + VERSION
}
+34
View File
@@ -255,3 +255,37 @@ class TRI3D_fill_mask():
image = to_torch_image(image).unsqueeze(0)
return (image,)
class TRI3D_is_only_trouser:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"pose_json_file": ("STRING", {
"default": "dwpose/keypoints"
})
}
}
RETURN_TYPES = ("BOOLEAN", )
FUNCTION = "main"
CATEGORY = "TRI3D"
def main(self, pose_json_file):
pose = json.load(open(pose_json_file))
height = pose['height']
width = pose['width']
keypoints = pose['keypoints']
points = [0,14,15,16,17,2,1,5]
point_to_part = {0:'nose',14:"left eye",15:"right eye",16:"left ear",17:"right ear",2:"left shoulder",1:"neck",5:"right shoulder"}
all_negative = True #if all face and shoulder points are negative means it is a bottom shot
for point in points:
x,y = keypoints[point]
if x > 0 and y > 0:
all_negative = False
print(f"{point_to_part[point]} exist")
return (all_negative,)