From ff2cd75e6fe467033e1b81e8d49ca6ec36f528da Mon Sep 17 00:00:00 2001
From: Ubuntu
 <ubuntu@researcha108.j1bjfcuzwdvepdrmknvfc3whuc.cx.internal.cloudapp.net>
Date: Thu, 22 Aug 2024 12:55:16 +0000
Subject: [PATCH] added node to detect trouser images

---
 __init__.py      | 10 ++++++----
 utility_nodes.py | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/__init__.py b/__init__.py
index a7eff72..264997c 100644
--- a/__init__.py
+++ b/__init__.py
@@ -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
 }
diff --git a/utility_nodes.py b/utility_nodes.py
index 49fea43..cbcb6b4 100644
--- a/utility_nodes.py
+++ b/utility_nodes.py
@@ -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,)
\ No newline at end of file