Merge pull request #60 from TRI3D-LC/saveflattened

flattened

.gitignore

@@ -5,10 +5,11 @@ venv
 .DS_Store
 checkpoints/
 checkpoint/
-ckpt/
 .env
 .pth
 cloth-segmentation/model/cloth_segm.pth
 
 dwpose/keypoints/
 huggingface/
+
+safetychecker/model.safetensors

__init__.py

@@ -1,5 +1,4 @@
 # v1.1.0
-
 import os
 import os.path
 tri3d_custom_nodes_path = os.path.dirname(os.path.abspath(__file__))
@@ -10,22 +9,29 @@ import torch.nn.functional as F
 import hashlib
 import comfy.model_management as model_management
 import folder_paths
-from PIL import Image, ImageOps, ImageDraw
+from PIL import Image, ImageOps
 sys.path.append(tri3d_custom_nodes_path)
-sys.path.append('.')
-sys.path.append('..')
 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, TRI3D_is_only_trouser, TRI3D_extract_facer_mask
+from utility_nodes import TRI3D_extract_facer_mask
 
 from .AEMatter import (load_AEMatter_Model, run_AEMatter_inference)
 from .light_layer import main_light_layer
+from .remove_small_mask_islands import TRI3D_RemoveSmallMaskIslands
+from .mask_area_percentage import TRI3D_MaskAreaPercentage
+
+
+from .image_stack import (
+    H_Stack_Images,
+    SaveImage_absolute,
+    SaveText_absolute,
+    Wait_And_Read_File,
+)
+
 
-from torchvision.transforms import ToPILImage, ToTensor
-from idm_masking.humanparsing.run_parsing import Parsing
-from idm_masking.openpose.run_openpose import OpenPose
 
 def from_torch_image(image):
     image = image.squeeze().cpu().numpy() * 255.0
@@ -1926,7 +1932,7 @@ class TRI3DDWPose_Preprocessor:
             cur_file_dir = os.path.dirname(os.path.realpath(__file__))
             save_file_path = os.path.join(cur_file_dir,
                                           filename_path)
-            json.dump(pose_dict, open(save_file_path, 'w'))
+            json.dump(pose_dict, open(save_file_path, 'w'), indent=4)
             np_result = cv2.resize(np_result, (W, H),
                                    interpolation=cv2.INTER_AREA)
             out_image_list.append(
@@ -1937,6 +1943,8 @@ class TRI3DDWPose_Preprocessor:
         out_image = torch.stack(out_image_list, dim=0)
         del model
 
+        # print(save_file_path, "save_file_path")
+
         return (out_image, save_file_path)
 
 
@@ -2662,9 +2670,12 @@ class TRI3D_reLUM:
                 mu_2, sigma_2 = get_mu_sigma(array_input=image_2[:, :, i],
                                              mask_input=mask_2)
 
+                sigma_calculated = sigma_1 * factor_sigma[i]
+                if factor_sigma[i] < 0:
+                    sigma_calculated = sigma_2
                 image_2[:, :, i] = (
                     ((image_2[:, :, i] - mu_2) / sigma_2) *
-                    (sigma_1 * factor_sigma[i])) + (mu_1 * factor_mean[i])
+                    sigma_calculated) + (mu_1 * factor_mean[i])
 
             image_2 = np.clip(image_2, 0, 255)
             image_2 = image_2.astype(dtype=np.uint8)
@@ -2839,9 +2850,14 @@ class TRI3D_recolor_LAB:
                 mu_2, sigma_2 = get_mu_sigma(array_input=image_2[:, :, i],
                                              mask_input=mask_2)
 
+                sigma_calculated = sigma_1 * factor_sigma[i]
+                if factor_sigma[i] < 0:
+                    sigma_calculated = sigma_2
+
                 image_2[:, :, i] = (
                     ((image_2[:, :, i] - mu_2) / sigma_2) *
-                    (sigma_1 * factor_sigma[i])) + (mu_1 * factor_mean[i])
+                    sigma_calculated) + (mu_1 * factor_mean[i])
+
 
             image_2 = np.clip(image_2, 0, 255)
             image_2 = image_2.astype(dtype=np.uint8)
@@ -3672,270 +3688,19 @@ class TRI3D_BGREMOVE_MEGA():
 
             batch_results = torch.stack(batch_results)
             batch_results_masks = torch.stack(batch_results_masks)
-
             return (batch_results,batch_results_masks)
                 
 
-#####
 
 
 
-label_map = {
-    "background": 0,
-    "hat": 1,
-    "hair": 2,
-    "sunglasses": 3,
-    "upper_clothes": 4,
-    "skirt": 5,
-    "pants": 6,
-    "dress": 7,
-    "belt": 8,
-    "left_shoe": 9,
-    "right_shoe": 10,
-    "head": 11,
-    "left_leg": 12,
-    "right_leg": 13,
-    "left_arm": 14,
-    "right_arm": 15,
-    "bag": 16,
-    "scarf": 17,
-}
-
-full_body_garments = [
-    "jumpsuit", "romper", "onesie", "coverall", "catsuit", "one-peice",
-    "overalls", "dungarees", "boilersuit", "suit", "wetsuit",
-    "swimsuit", "spacesuit", "leotard", "unitard", "kimono",
-    "caftan", "abaya", "sari", "dashiki", "djellaba", "muumuu", "full_body"
-]
-
-upper_body_garments = [
-    "shirt", "blouse", "t-shirt", "sweater", "cardigan", "jacket",
-    "blazer", "hoodie", "tank top", "vest", "crop top", "tunic", "bra",
-    "polo shirt", "sweatshirt", "pullover", "turtleneck", "halter top",
-    "bolero", "poncho", "shrug", "camisole", "bustier", "corset",
-    "coat", "parka", "windbreaker", "upper_body"
-]
-
-lower_body_garments = [
-    "pants", "trousers", "jeans", "shorts", "skirt", "leggings",
-    "joggers", "sweatpants", "chinos", "khakis", "cargo pants",
-    "culottes", "capris", "palazzo", "bermuda shorts", "panties",
-    "mini skirt", "midi skirt", "maxi skirt", "sarong", "kilt",
-    "dhoti", "harem pants", "bloomers", "lower_body"
-]
-
-
-MAX_RESOLUTION = 16384
-
-class TRI3DBodyMask:
-    
-    def __init__(self):
-        self.parsing_model = Parsing(0)
-        self.openpose_model = OpenPose(0)
-    
-    
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-                "required": {
-                "human_img": ("IMAGE",),
-                "garment_description": ("STRING", {"multiline": True}),
-                "width": ("INT", {"default": 768, "min": 0, "max": MAX_RESOLUTION}),
-                "height": ("INT", {"default": 1024, "min": 0, "max": MAX_RESOLUTION}),
-                "dilate": ("INT", {"default": 1, "min": -50, "max": 50}),
-            }
-        }
-
-    RETURN_TYPES = ("IMAGE","MASK")
-    FUNCTION = "body_mask"
-    CATEGORY = "ComfyUI-IDM-VTON"
-    
-
-
-    def extend_arm_mask(self, wrist, elbow, scale):
-        wrist = elbow + scale * (wrist - elbow)
-        return wrist
-
-    def hole_fill(self, img):
-        img = np.pad(img[1:-1, 1:-1], pad_width = 1, mode = 'constant', constant_values=0)
-        img_copy = img.copy()
-        mask = np.zeros((img.shape[0] + 2, img.shape[1] + 2), dtype=np.uint8)
-
-        cv2.floodFill(img, mask, (0, 0), 255)
-        img_inverse = cv2.bitwise_not(img)
-        dst = cv2.bitwise_or(img_copy, img_inverse)
-        return dst
-
-    def refine_mask(self, mask):
-        contours, hierarchy = cv2.findContours(mask.astype(np.uint8),
-                                            cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_L1)
-        area = []
-        for j in range(len(contours)):
-            a_d = cv2.contourArea(contours[j], True)
-            area.append(abs(a_d))
-        refined_mask = np.zeros_like(mask).astype(np.uint8)
-        if len(area) != 0:
-            i = area.index(max(area))
-            cv2.drawContours(refined_mask, contours, i, color=255, thickness=-1)
-
-        return refined_mask
-    
-    def check_category(self, word_list, input_string):
-        input_string = input_string.lower()
-        for word in word_list:
-            if word.lower() in input_string:
-                return True
-        return False
-
-    def body_mask(self, human_img: Image.Image, garment_description, width, height, dilate):
-        human_img =  torch.transpose(human_img, 0, -1)
-        human_img = torch.squeeze(human_img, dim=-1)
-        
-        human_img = ToPILImage()(human_img)
-        human_img = human_img.resize((width,height))
-        keypoint = self.openpose_model(human_img.resize((384, 512)))
-        model_img, _ = self.parsing_model(human_img.resize((384, 512)))
-        im_parse = model_img.resize((width, height), Image.NEAREST)
-        parse_array = np.array(im_parse)
-        
-        f = False
-        u = False
-        l = False
-        
-        category = "dresses"
-        if self.check_category(full_body_garments, garment_description):
-            f = True
-        if self.check_category(upper_body_garments, garment_description):
-            u = True
-        if self.check_category(lower_body_garments, garment_description):
-            l = True
-
-        if f:
-            category = "dresses"
-        elif u and l:
-            category = "dresses"
-        elif u:
-            category = "upper_body"
-        elif l:
-            category = "lower_body"        
-        
-        arm_width = 45
-        
-        parse_head = (parse_array == 1).astype(np.float32) + \
-                    (parse_array == 3).astype(np.float32) + \
-                    (parse_array == 11).astype(np.float32)
-
-        parser_mask_fixed = (parse_array == label_map["left_shoe"]).astype(np.float32) + \
-                            (parse_array == label_map["right_shoe"]).astype(np.float32) + \
-                            (parse_array == label_map["hat"]).astype(np.float32) + \
-                            (parse_array == label_map["sunglasses"]).astype(np.float32) + \
-                            (parse_array == label_map["bag"]).astype(np.float32)
-
-        parser_mask_changeable = (parse_array == label_map["background"]).astype(np.float32)
-
-        arms_left = (parse_array == 14).astype(np.float32)
-        arms_right = (parse_array == 15).astype(np.float32)
-
-        if category == 'dresses':
-            parse_mask = (parse_array == 7).astype(np.float32) + \
-                        (parse_array == 4).astype(np.float32) + \
-                        (parse_array == 5).astype(np.float32) + \
-                        (parse_array == 6).astype(np.float32)
-            parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
-
-        elif category == 'upper_body':
-            parse_mask = (parse_array == 4).astype(np.float32) + (parse_array == 7).astype(np.float32)
-            parser_mask_fixed_lower_cloth = (parse_array == label_map["skirt"]).astype(np.float32) + \
-                                            (parse_array == label_map["pants"]).astype(np.float32)
-            parser_mask_fixed += parser_mask_fixed_lower_cloth
-            parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
-
-        elif category == 'lower_body':
-            parse_mask = (parse_array == 6).astype(np.float32) + \
-                        (parse_array == 12).astype(np.float32) + \
-                        (parse_array == 13).astype(np.float32) + \
-                        (parse_array == 5).astype(np.float32)
-            parser_mask_fixed += (parse_array == label_map["upper_clothes"]).astype(np.float32) + \
-                                (parse_array == 14).astype(np.float32) + \
-                                (parse_array == 15).astype(np.float32)
-            parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
-        else:
-            raise NotImplementedError
-
-        # Load pose points
-        pose_data = keypoint["pose_keypoints_2d"]
-        pose_data = np.array(pose_data)
-        pose_data = pose_data.reshape((-1, 2))
-
-        im_arms_left = Image.new('L', (width, height))
-        im_arms_right = Image.new('L', (width, height))
-        arms_draw_left = ImageDraw.Draw(im_arms_left)
-        arms_draw_right = ImageDraw.Draw(im_arms_right)
-        if category == 'dresses' or category == 'upper_body':
-            shoulder_right = np.multiply(tuple(pose_data[2][:2]), height / 512.0)
-            shoulder_left = np.multiply(tuple(pose_data[5][:2]), height / 512.0)
-            elbow_right = np.multiply(tuple(pose_data[3][:2]), height / 512.0)
-            elbow_left = np.multiply(tuple(pose_data[6][:2]), height / 512.0)
-            wrist_right = np.multiply(tuple(pose_data[4][:2]), height / 512.0)
-            wrist_left = np.multiply(tuple(pose_data[7][:2]), height / 512.0)
-            ARM_LINE_WIDTH = int(arm_width / 512 * height)
-            size_left = [shoulder_left[0] - ARM_LINE_WIDTH // 2, shoulder_left[1] - ARM_LINE_WIDTH // 2, shoulder_left[0] + ARM_LINE_WIDTH // 2, shoulder_left[1] + ARM_LINE_WIDTH // 2]
-            size_right = [shoulder_right[0] - ARM_LINE_WIDTH // 2, shoulder_right[1] - ARM_LINE_WIDTH // 2, shoulder_right[0] + ARM_LINE_WIDTH // 2,
-                        shoulder_right[1] + ARM_LINE_WIDTH // 2]
-            
-
-            if wrist_right[0] <= 1. and wrist_right[1] <= 1.:
-                im_arms_right = arms_right
-            else:
-                wrist_right = self.extend_arm_mask(wrist_right, elbow_right, 1.2)
-                arms_draw_right.line(np.concatenate((shoulder_right, elbow_right, wrist_right)).astype(np.uint16).tolist(), 'white', ARM_LINE_WIDTH, 'curve')
-                arms_draw_right.arc(size_right, 0, 360, 'white', ARM_LINE_WIDTH // 2)
-
-            if wrist_left[0] <= 1. and wrist_left[1] <= 1.:
-                im_arms_left = arms_left
-            else:
-                wrist_left = self.extend_arm_mask(wrist_left, elbow_left, 1.2)
-                arms_draw_left.line(np.concatenate((wrist_left, elbow_left, shoulder_left)).astype(np.uint16).tolist(), 'white', ARM_LINE_WIDTH, 'curve')
-                arms_draw_left.arc(size_left, 0, 360, 'white', ARM_LINE_WIDTH // 2)
-
-            hands_left = np.logical_and(np.logical_not(im_arms_left), arms_left)
-            hands_right = np.logical_and(np.logical_not(im_arms_right), arms_right)
-            parser_mask_fixed += hands_left + hands_right
-
-        parser_mask_fixed = np.logical_or(parser_mask_fixed, parse_head)
-        parse_mask = cv2.dilate(parse_mask, np.ones((5, 5), np.uint16), iterations=5)
-        if category == 'dresses' or category == 'upper_body':
-            neck_mask = (parse_array == 18).astype(np.float32)
-            neck_mask = cv2.dilate(neck_mask, np.ones((5, 5), np.uint16), iterations=1)
-            neck_mask = np.logical_and(neck_mask, np.logical_not(parse_head))
-            parse_mask = np.logical_or(parse_mask, neck_mask)
-            arm_mask = cv2.dilate(np.logical_or(im_arms_left, im_arms_right).astype('float32'), np.ones((5, 5), np.uint16), iterations=4)
-            parse_mask += np.logical_or(parse_mask, arm_mask)
-
-        parse_mask = np.logical_and(parser_mask_changeable, np.logical_not(parse_mask))
-
-        parse_mask_total = np.logical_or(parse_mask, parser_mask_fixed)
-        inpaint_mask = 1 - parse_mask_total
-        img = np.where(inpaint_mask, 255, 0)
-        dst = self.hole_fill(img.astype(np.uint8))
-        dst = self.refine_mask(dst)
-        mask = (dst / 255 * 1).astype(np.uint8)
-        
-        kernel = np.ones((3, 3), np.uint8)
-        if dilate >= 0:
-            mask = cv2.dilate(mask, kernel, iterations=int(dilate))
-        else:
-            mask = cv2.erode(mask, kernel, iterations=int(abs(dilate)))
-
-        mask = Image.fromarray(mask * 255)
-        mask = ToTensor()(mask)
-        
-        return (mask, mask,)
-
-
-#####
 
 from photoroom import TRI3D_photoroom_bgremove_api
+from smart_box import TRI3D_SmartBox, TRI3D_Skip_HeadMask, TRI3D_Skip_HeadMask_AddNeck, TRI3D_Image_extend, TRI3D_Smart_Depth, TRI3D_NarrowfyImage, TRI3D_Skip_LipMask
+from nsfw import TRI3DNSFWFilter
+from cut_by_mask_aspect_ratio import TRI3D_CutByMaskAspectRatio
+from string_check import TRI3D_StringContains
+from .dwpose_conversion import SaveFlattenedPoseKpsAsJsonFile
 
 # A dictionary that contains all nodes you want to export with their names
 # NOTE: names should be globally unique
@@ -3986,11 +3751,35 @@ NODE_CLASS_MAPPINGS = {
     'tri3d-run_AEMatter_inference': run_AEMatter_inference,
     "tri3d-bgremove-mega" :TRI3D_BGREMOVE_MEGA,
     'tri3d-flexible_color_extract' : main_light_layer,
-    "tri3d-automatic-body-mask": TRI3DBodyMask,
+    '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_is_only_trouser": TRI3D_is_only_trouser,
+    "tri3d_extract_facer_mask":TRI3D_extract_facer_mask,
+    "tri3d_H_Stack_Images": H_Stack_Images,
+    "tri3d_SaveImage_absolute":SaveImage_absolute,
+    "tri3d_SaveText_absolute":SaveText_absolute,
+    "tri3d_Wait_And_Read_File":Wait_And_Read_File,
+    "tri3d_SmartBox": TRI3D_SmartBox,
+    "tri3d_Skip_HeadMask": TRI3D_Skip_HeadMask,
+    "tri3d_Skip_HeadMask_AddNeck": TRI3D_Skip_HeadMask_AddNeck,
+    "tri3d_Skip_LipMask": TRI3D_Skip_LipMask,
+    "tri3d_Image_extend": TRI3D_Image_extend,
+    "tri3d_Smart_Depth": TRI3D_Smart_Depth,
+    "tri3d_NSFWFilter": TRI3DNSFWFilter,
+    "tri3d_NarrowfyImage": TRI3D_NarrowfyImage,
+    "tri3d_Skip_LipMask": TRI3D_Skip_LipMask,
+    "tri3d_Remove_Small_Mask_Islands": TRI3D_RemoveSmallMaskIslands,
+    "tri3d_CutByMaskAspectRatio": TRI3D_CutByMaskAspectRatio,
+    "tri3d_StringContains": TRI3D_StringContains,
+    "tri3d_MaskAreaPercentage": TRI3D_MaskAreaPercentage,
+    "tri3d_SaveFlattenedPoseKpsAsJsonFile": SaveFlattenedPoseKpsAsJsonFile,
+
 }
 
 
-VERSION = "4.4"
+VERSION = "5.1.0"
 # 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,
@@ -4041,5 +3830,27 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     'tri3d-run_AEMatter_inference': 'Run AEMatter inference' + ' v' + VERSION,
     "tri3d-bgremove-mega": "BG Remove Mega" + " v" + VERSION,
     'tri3d-flexible_color_extract': "Flexible color extract" + " v" + VERSION,
-    "tri3d-automatic-body-mask": "Find correct Body Mask from Flats" + " v" + VERSION,
+    '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_is_only_trouser": "Is only trouser" + " v" + VERSION,
+    "tri3d_extract_facer_mask": "Extract facer mask" + " v" + VERSION,
+    "tri3d_H_Stack_Images": "Stack images for cat vton with flux" + " v" + VERSION,
+    "tri3d_SaveImage_absolute": "Save image to an absolute path and provide text optional to control execution order" + " v" + VERSION,
+    "tri3d_SaveText_absolute": "Save text to an absolute path and provide text optional to control execution order " + " v" + VERSION,
+    "tri3d_Wait_And_Read_File": "Wait and read text file, optional control from text " + " v" + VERSION,
+    "tri3d_SmartBox": "Smart Box" + " v" + VERSION,
+    "tri3d_Skip_HeadMask": "Skip Head Mask" + " v" + VERSION,
+    "tri3d_Skip_HeadMask_AddNeck": "Skip Head Mask and add neck" + " v" + VERSION,
+    "tri3d_Skip_LipMask": "Skip Lip Mask" + " v" + VERSION,
+    "tri3d_NSFWFilter": "TRI3D NSFW Filter" + " v" + VERSION,
+    "tri3d_Image_extend": "Image extend" + " v" + VERSION,
+    "tri3d_Smart_Depth": "Smart Depth" + " v" + VERSION,
+    "tri3d_NarrowfyImage": "Narrowfy Image" + " v" + VERSION,
+    "tri3d_Remove_Small_Mask_Islands": "Remove Small Mask Islands" + " v" + VERSION,
+    "tri3d_CutByMaskAspectRatio": "Cut by mask aspect ratio" + " v" + VERSION,
+    "tri3d_StringContains": "String contains" + " v" + VERSION,
+    "tri3d_MaskAreaPercentage": "Mask Area Percentage" + " v" + VERSION,
+    "tri3d_SaveFlattenedPoseKpsAsJsonFile": "Save Flattened Pose Keypoints as JSON File" + " v" + VERSION,
 }

cut_by_mask_aspect_ratio.py

@@ -0,0 +1,183 @@
+import os
+import cv2
+import numpy as np
+import torch
+
+class TRI3D_CutByMaskAspectRatio:
+    """
+    ComfyUI node that crops an image based on a mask's bounding box,
+    adjusts the aspect ratio, and resizes to specified dimensions.
+    """
+    
+    def from_torch_image(self, image):
+        """Convert a torch tensor image to numpy array for OpenCV processing"""
+        image = image.cpu().numpy() * 255.0
+        image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    def to_torch_image(self, image):
+        """Convert numpy array back to torch tensor format"""
+        image = image.astype(dtype=np.float32)
+        image /= 255.0
+        image = torch.from_numpy(image)
+        return image
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "mask": ("IMAGE",),
+                "margin": ("INT", {"default": 10, "min": 0, "max": 100, "step": 1}),
+                "target_width": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 8}),
+                "target_height": ("INT", {"default": 512, "min": 64, "max": 4096, "step": 8}),
+                "padding_color": ("INT", {"default": 255, "min": 0, "max": 255, "step": 1}),
+            },
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE",)
+    CATEGORY = "TRI3D"
+
+    def run(self, image, mask, margin, target_width, target_height, padding_color=255):
+        # Convert Torch images to OpenCV format
+        cv_image = self.from_torch_image(image)
+        cv_mask = self.from_torch_image(mask)
+        
+        # Remove batch dimension if present
+        if len(cv_image.shape) == 4:
+            cv_image = cv_image[0]
+        if len(cv_mask.shape) == 4:
+            cv_mask = cv_mask[0]
+        
+        # Convert mask to grayscale if it's not already
+        if len(cv_mask.shape) == 3 and cv_mask.shape[2] > 1:
+            mask_gray = cv2.cvtColor(cv_mask, cv2.COLOR_RGB2GRAY)
+        else:
+            mask_gray = cv_mask[:, :, 0]
+        
+        # Create binary mask
+        _, binary_mask = cv2.threshold(mask_gray, 127, 255, cv2.THRESH_BINARY)
+        
+        # Find contours in the binary mask
+        contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+        if not contours:
+            # If no contours found, return the original image
+            print("No contours found in mask. Returning original image.")
+            return (image,)
+        
+        # Find bounding box around all contours
+        x_min, y_min = float('inf'), float('inf')
+        x_max, y_max = 0, 0
+        
+        for contour in contours:
+            x, y, w, h = cv2.boundingRect(contour)
+            x_min = min(x_min, x)
+            y_min = min(y_min, y)
+            x_max = max(x_max, x + w)
+            y_max = max(y_max, y + h)
+        
+        # Add margin to bounding box
+        x_min = max(0, x_min - margin)
+        y_min = max(0, y_min - margin)
+        x_max = min(cv_image.shape[1], x_max + margin)
+        y_max = min(cv_image.shape[0], y_max + margin)
+        
+        # Current dimensions of the bounding box
+        height = y_max - y_min
+        width = x_max - x_min
+        
+        # Calculate the target aspect ratio (width/height)
+        target_aspect_ratio = target_width / target_height
+        
+        # Calculate current aspect ratio
+        current_aspect_ratio = width / height
+        
+        # Adjust width to match the target aspect ratio while keeping height constant
+        if current_aspect_ratio < target_aspect_ratio:
+            # Current width is too narrow - need to extend it
+            # Calculate the required width for the target aspect ratio
+            required_width = int(height * target_aspect_ratio)
+            width_difference = required_width - width
+            
+            # Calculate how much to extend on each side
+            left_extend = width_difference // 2
+            right_extend = width_difference - left_extend
+            
+            # Calculate new potential boundaries
+            new_x_min = x_min - left_extend
+            new_x_max = x_max + right_extend
+            
+            # Check if the new boundaries are within the original image
+            left_padding_needed = abs(min(0, new_x_min))
+            right_padding_needed = max(0, new_x_max - cv_image.shape[1])
+            
+            # Adjust boundaries to be within the original image
+            new_x_min = max(0, new_x_min)
+            new_x_max = min(cv_image.shape[1], new_x_max)
+            
+            # Get the portion of the original image within valid boundaries
+            extended_image = cv_image[y_min:y_max, new_x_min:new_x_max]
+            
+            # If we need padding (i.e., extension goes beyond image boundaries)
+            if left_padding_needed > 0 or right_padding_needed > 0:
+                # Create canvas with padding color
+                num_channels = extended_image.shape[2] if len(extended_image.shape) == 3 else 1
+                if num_channels == 1:
+                    canvas = np.full((height, required_width), padding_color, dtype=np.uint8)
+                else:
+                    canvas = np.full((height, required_width, num_channels), padding_color, dtype=np.uint8)
+                
+                # Calculate the position to place the extended image
+                place_x = left_padding_needed
+                
+                # Place the extended image on the canvas
+                if num_channels == 1:
+                    canvas[:, place_x:place_x+extended_image.shape[1]] = extended_image
+                else:
+                    canvas[:, place_x:place_x+extended_image.shape[1], :] = extended_image
+                
+                # Use the canvas as our cropped image
+                cropped_image = canvas
+            else:
+                # No padding needed, use the extended image
+                cropped_image = extended_image
+            
+        elif current_aspect_ratio > target_aspect_ratio:
+            # Current width is too wide, crop it
+            new_width = int(height * target_aspect_ratio)
+            width_difference = width - new_width
+            
+            # Crop equally from both sides if possible
+            left_crop = width_difference // 2
+            right_crop = width_difference - left_crop
+            
+            # Apply the crop
+            cropped_image = cv_image[y_min:y_max, x_min+left_crop:x_max-right_crop]
+        else:
+            # Aspect ratio is already correct
+            cropped_image = cv_image[y_min:y_max, x_min:x_max]
+        
+        # Resize the cropped/padded image to the target dimensions using Lanczos interpolation
+        resized_image = cv2.resize(cropped_image, (target_width, target_height), interpolation=cv2.INTER_LANCZOS4)
+        
+        # Convert back to torch format
+        torch_image = self.to_torch_image(resized_image)
+        
+        # Add batch dimension back
+        torch_image = torch_image.unsqueeze(0)
+        
+        return (torch_image,)
+
+# Node registration for ComfyUI
+NODE_CLASS_MAPPINGS = {
+    "TRI3D_CutByMaskAspectRatio": TRI3D_CutByMaskAspectRatio
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "TRI3D_CutByMaskAspectRatio": "TRI3D Cut By Mask Aspect Ratio"
+}

download_idm_checkpoint.sh

@@ -1,22 +0,0 @@
-#!/bin/bash
-
-# Create main checkpoint directory
-mkdir -p ckpt
-
-# Create subdirectories
-mkdir -p ckpt/densepose
-mkdir -p ckpt/humanparsing
-mkdir -p ckpt/openpose/ckpts
-
-# Download files for densepose
-wget -P ckpt/densepose https://huggingface.co/spaces/yisol/IDM-VTON/resolve/main/ckpt/densepose/model_final_162be9.pkl
-
-# Download files for humanparsing
-wget -P ckpt/humanparsing https://huggingface.co/spaces/yisol/IDM-VTON/resolve/main/ckpt/humanparsing/parsing_atr.onnx
-wget -P ckpt/humanparsing https://huggingface.co/spaces/yisol/IDM-VTON/resolve/main/ckpt/humanparsing/parsing_lip.onnx
-
-# Download files for openpose
-wget -P ckpt/openpose/ckpts https://huggingface.co/spaces/yisol/IDM-VTON/resolve/main/ckpt/openpose/ckpts/body_pose_model.pth
-
-echo "Download completed!"
-

dwpose/comfy_utils.py

@@ -274,4 +274,6 @@ def switch_to_backpose(input_keypoints, input_width):
         x,y = input_keypoints[i]
         input_keypoints[i] = [input_width - x, y]
 
-    return input_keypoints
+    return input_keypoints
+
+

dwpose_conversion.py

@@ -0,0 +1,111 @@
+import os
+import json
+import torch
+import numpy as np
+import folder_paths
+
+print("Loading TRI3D_SavePoseKeypointsJSON module")
+
+class SaveFlattenedPoseKpsAsJsonFile:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "pose_kps": ("POSE_KEYPOINT",),
+                "file_path": ("STRING", {"default": "dwpose/keypoints/input.json"})
+            }
+        }
+    RETURN_TYPES = (
+        "STRING",
+    )
+    FUNCTION = "save_flattened_pose_kps"
+    OUTPUT_NODE = True
+    CATEGORY = "ControlNet Preprocessors/Pose Keypoint Postprocess"
+
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+        self.type = "output"
+        self.prefix_append = ""
+
+    def _flatten_openpose_dict(self, pose_dict: dict) -> dict:
+        """
+        Converts a single OpenPose dictionary into flattened format.
+        """
+        # Get canvas dimensions from the input dictionary
+        H = pose_dict.get('canvas_height', 512)
+        W = pose_dict.get('canvas_width', 512)
+        
+        flat_keypoints = []
+
+        # Check if any person was detected
+        if not pose_dict.get('people'):
+            # If no people, return a list of 130 invalid keypoints
+            flat_keypoints.extend([[-1, -1]] * 130)
+            return {"height": H, "width": W, "keypoints": flat_keypoints}
+
+        person = pose_dict['people'][0]  # Process the first person found
+
+        # Helper function to process each body part
+        def process_part(keypoints_data, expected_length):
+            processed_kps = []
+            if keypoints_data:
+                # Iterate in steps of 3 (x, y, confidence)
+                for i in range(0, len(keypoints_data), 3):
+                    x, y, conf = keypoints_data[i], keypoints_data[i+1], keypoints_data[i+2]
+                    # Use confidence score to check for validity. If 0, it's a missing point.
+                    if conf > 0:
+                        processed_kps.append([x, y])
+                    else:
+                        processed_kps.append([-1, -1])
+            
+            # Ensure the list has the exact expected length
+            while len(processed_kps) < expected_length:
+                processed_kps.append([-1, -1])
+            
+            return processed_kps
+
+        # Process parts in order: body -> face -> left hand -> right hand
+        body_kps = process_part(person.get('pose_keypoints_2d'), 18)
+        face_kps = process_part(person.get('face_keypoints_2d'), 70)
+        left_hand_kps = process_part(person.get('hand_left_keypoints_2d'), 21)
+        right_hand_kps = process_part(person.get('hand_right_keypoints_2d'), 21)
+
+        # Combine all parts into the final flat list
+        flat_keypoints.extend(body_kps)
+        flat_keypoints.extend(face_kps)
+        flat_keypoints.extend(left_hand_kps)
+        flat_keypoints.extend(right_hand_kps)
+        
+        return {"height": H, "width": W, "keypoints": flat_keypoints}
+
+    def save_flattened_pose_kps(self, pose_kps, file_path):
+        # filename_prefix += self.prefix_append
+        
+        # # Get the save path using the first pose keypoint's dimensions
+        # full_output_folder, filename, counter, subfolder, filename_prefix = \
+        #     folder_paths.get_save_image_path(filename_prefix, self.output_dir, 
+        #                                   pose_kps[0]["canvas_width"], 
+        #                                   pose_kps[0]["canvas_height"])
+        
+        # Process each pose keypoint in the batch
+        flattened_poses = []
+        for pose_dict in pose_kps:
+            flattened_data = self._flatten_openpose_dict(pose_dict)
+            flattened_poses.append(flattened_data)
+            
+        # # Save the flattened data
+        # file = f"{filename}_{counter:05}.json"
+        # save_path = os.path.join(full_output_folder, file)
+
+        cur_file_dir = os.path.dirname(os.path.realpath(__file__))
+        save_path = os.path.join(cur_file_dir,
+                                          file_path)
+        
+        with open(save_path, 'w') as f:
+            if len(flattened_poses) == 1:
+                json.dump(flattened_poses[0], f, indent=4)  # Save single pose directly
+            else:
+                json.dump(flattened_poses, f, indent=4)  # Save batch as array
+                
+        print(f"Saved flattened pose keypoints to: {save_path}")
+        return (save_path,)

idm_masking/humanparsing/modules/__init__.py

@@ -1,5 +0,0 @@
-from .bn import ABN, InPlaceABN, InPlaceABNSync
-from .functions import ACT_RELU, ACT_LEAKY_RELU, ACT_ELU, ACT_NONE
-from .misc import GlobalAvgPool2d, SingleGPU
-from .residual import IdentityResidualBlock
-from .dense import DenseModule

idm_masking/humanparsing/modules/bn.py

@@ -1,132 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as functional
-
-try:
-    from queue import Queue
-except ImportError:
-    from Queue import Queue
-
-from .functions import *
-
-
-class ABN(nn.Module):
-    """Activated Batch Normalization
-
-    This gathers a `BatchNorm2d` and an activation function in a single module
-    """
-
-    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, activation="leaky_relu", slope=0.01):
-        """Creates an Activated Batch Normalization module
-
-        Parameters
-        ----------
-        num_features : int
-            Number of feature channels in the input and output.
-        eps : float
-            Small constant to prevent numerical issues.
-        momentum : float
-            Momentum factor applied to compute running statistics as.
-        affine : bool
-            If `True` apply learned scale and shift transformation after normalization.
-        activation : str
-            Name of the activation functions, one of: `leaky_relu`, `elu` or `none`.
-        slope : float
-            Negative slope for the `leaky_relu` activation.
-        """
-        super(ABN, self).__init__()
-        self.num_features = num_features
-        self.affine = affine
-        self.eps = eps
-        self.momentum = momentum
-        self.activation = activation
-        self.slope = slope
-        if self.affine:
-            self.weight = nn.Parameter(torch.ones(num_features))
-            self.bias = nn.Parameter(torch.zeros(num_features))
-        else:
-            self.register_parameter('weight', None)
-            self.register_parameter('bias', None)
-        self.register_buffer('running_mean', torch.zeros(num_features))
-        self.register_buffer('running_var', torch.ones(num_features))
-        self.reset_parameters()
-
-    def reset_parameters(self):
-        nn.init.constant_(self.running_mean, 0)
-        nn.init.constant_(self.running_var, 1)
-        if self.affine:
-            nn.init.constant_(self.weight, 1)
-            nn.init.constant_(self.bias, 0)
-
-    def forward(self, x):
-        x = functional.batch_norm(x, self.running_mean, self.running_var, self.weight, self.bias,
-                                  self.training, self.momentum, self.eps)
-
-        if self.activation == ACT_RELU:
-            return functional.relu(x, inplace=True)
-        elif self.activation == ACT_LEAKY_RELU:
-            return functional.leaky_relu(x, negative_slope=self.slope, inplace=True)
-        elif self.activation == ACT_ELU:
-            return functional.elu(x, inplace=True)
-        else:
-            return x
-
-    def __repr__(self):
-        rep = '{name}({num_features}, eps={eps}, momentum={momentum},' \
-              ' affine={affine}, activation={activation}'
-        if self.activation == "leaky_relu":
-            rep += ', slope={slope})'
-        else:
-            rep += ')'
-        return rep.format(name=self.__class__.__name__, **self.__dict__)
-
-
-class InPlaceABN(ABN):
-    """InPlace Activated Batch Normalization"""
-
-    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, activation="leaky_relu", slope=0.01):
-        """Creates an InPlace Activated Batch Normalization module
-
-        Parameters
-        ----------
-        num_features : int
-            Number of feature channels in the input and output.
-        eps : float
-            Small constant to prevent numerical issues.
-        momentum : float
-            Momentum factor applied to compute running statistics as.
-        affine : bool
-            If `True` apply learned scale and shift transformation after normalization.
-        activation : str
-            Name of the activation functions, one of: `leaky_relu`, `elu` or `none`.
-        slope : float
-            Negative slope for the `leaky_relu` activation.
-        """
-        super(InPlaceABN, self).__init__(num_features, eps, momentum, affine, activation, slope)
-
-    def forward(self, x):
-        x, _, _ = inplace_abn(x, self.weight, self.bias, self.running_mean, self.running_var,
-                           self.training, self.momentum, self.eps, self.activation, self.slope)
-        return x
-
-
-class InPlaceABNSync(ABN):
-    """InPlace Activated Batch Normalization with cross-GPU synchronization
-    This assumes that it will be replicated across GPUs using the same mechanism as in `nn.DistributedDataParallel`.
-    """
-
-    def forward(self, x):
-        x, _, _ =  inplace_abn_sync(x, self.weight, self.bias, self.running_mean, self.running_var,
-                                   self.training, self.momentum, self.eps, self.activation, self.slope)
-        return x
-
-    def __repr__(self):
-        rep = '{name}({num_features}, eps={eps}, momentum={momentum},' \
-              ' affine={affine}, activation={activation}'
-        if self.activation == "leaky_relu":
-            rep += ', slope={slope})'
-        else:
-            rep += ')'
-        return rep.format(name=self.__class__.__name__, **self.__dict__)
-
-

idm_masking/humanparsing/modules/deeplab.py

@@ -1,84 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as functional
-
-from models._util import try_index
-from .bn import ABN
-
-
-class DeeplabV3(nn.Module):
-    def __init__(self,
-                 in_channels,
-                 out_channels,
-                 hidden_channels=256,
-                 dilations=(12, 24, 36),
-                 norm_act=ABN,
-                 pooling_size=None):
-        super(DeeplabV3, self).__init__()
-        self.pooling_size = pooling_size
-
-        self.map_convs = nn.ModuleList([
-            nn.Conv2d(in_channels, hidden_channels, 1, bias=False),
-            nn.Conv2d(in_channels, hidden_channels, 3, bias=False, dilation=dilations[0], padding=dilations[0]),
-            nn.Conv2d(in_channels, hidden_channels, 3, bias=False, dilation=dilations[1], padding=dilations[1]),
-            nn.Conv2d(in_channels, hidden_channels, 3, bias=False, dilation=dilations[2], padding=dilations[2])
-        ])
-        self.map_bn = norm_act(hidden_channels * 4)
-
-        self.global_pooling_conv = nn.Conv2d(in_channels, hidden_channels, 1, bias=False)
-        self.global_pooling_bn = norm_act(hidden_channels)
-
-        self.red_conv = nn.Conv2d(hidden_channels * 4, out_channels, 1, bias=False)
-        self.pool_red_conv = nn.Conv2d(hidden_channels, out_channels, 1, bias=False)
-        self.red_bn = norm_act(out_channels)
-
-        self.reset_parameters(self.map_bn.activation, self.map_bn.slope)
-
-    def reset_parameters(self, activation, slope):
-        gain = nn.init.calculate_gain(activation, slope)
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                nn.init.xavier_normal_(m.weight.data, gain)
-                if hasattr(m, "bias") and m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
-            elif isinstance(m, ABN):
-                if hasattr(m, "weight") and m.weight is not None:
-                    nn.init.constant_(m.weight, 1)
-                if hasattr(m, "bias") and m.bias is not None:
-                    nn.init.constant_(m.bias, 0)
-
-    def forward(self, x):
-        # Map convolutions
-        out = torch.cat([m(x) for m in self.map_convs], dim=1)
-        out = self.map_bn(out)
-        out = self.red_conv(out)
-
-        # Global pooling
-        pool = self._global_pooling(x)
-        pool = self.global_pooling_conv(pool)
-        pool = self.global_pooling_bn(pool)
-        pool = self.pool_red_conv(pool)
-        if self.training or self.pooling_size is None:
-            pool = pool.repeat(1, 1, x.size(2), x.size(3))
-
-        out += pool
-        out = self.red_bn(out)
-        return out
-
-    def _global_pooling(self, x):
-        if self.training or self.pooling_size is None:
-            pool = x.view(x.size(0), x.size(1), -1).mean(dim=-1)
-            pool = pool.view(x.size(0), x.size(1), 1, 1)
-        else:
-            pooling_size = (min(try_index(self.pooling_size, 0), x.shape[2]),
-                            min(try_index(self.pooling_size, 1), x.shape[3]))
-            padding = (
-                (pooling_size[1] - 1) // 2,
-                (pooling_size[1] - 1) // 2 if pooling_size[1] % 2 == 1 else (pooling_size[1] - 1) // 2 + 1,
-                (pooling_size[0] - 1) // 2,
-                (pooling_size[0] - 1) // 2 if pooling_size[0] % 2 == 1 else (pooling_size[0] - 1) // 2 + 1
-            )
-
-            pool = functional.avg_pool2d(x, pooling_size, stride=1)
-            pool = functional.pad(pool, pad=padding, mode="replicate")
-        return pool

idm_masking/humanparsing/modules/dense.py

@@ -1,42 +0,0 @@
-from collections import OrderedDict
-
-import torch
-import torch.nn as nn
-
-from .bn import ABN
-
-
-class DenseModule(nn.Module):
-    def __init__(self, in_channels, growth, layers, bottleneck_factor=4, norm_act=ABN, dilation=1):
-        super(DenseModule, self).__init__()
-        self.in_channels = in_channels
-        self.growth = growth
-        self.layers = layers
-
-        self.convs1 = nn.ModuleList()
-        self.convs3 = nn.ModuleList()
-        for i in range(self.layers):
-            self.convs1.append(nn.Sequential(OrderedDict([
-                ("bn", norm_act(in_channels)),
-                ("conv", nn.Conv2d(in_channels, self.growth * bottleneck_factor, 1, bias=False))
-            ])))
-            self.convs3.append(nn.Sequential(OrderedDict([
-                ("bn", norm_act(self.growth * bottleneck_factor)),
-                ("conv", nn.Conv2d(self.growth * bottleneck_factor, self.growth, 3, padding=dilation, bias=False,
-                                   dilation=dilation))
-            ])))
-            in_channels += self.growth
-
-    @property
-    def out_channels(self):
-        return self.in_channels + self.growth * self.layers
-
-    def forward(self, x):
-        inputs = [x]
-        for i in range(self.layers):
-            x = torch.cat(inputs, dim=1)
-            x = self.convs1[i](x)
-            x = self.convs3[i](x)
-            inputs += [x]
-
-        return torch.cat(inputs, dim=1)

idm_masking/humanparsing/modules/functions.py

@@ -1,245 +0,0 @@
-import pdb
-from os import path
-import torch
-import torch.distributed as dist
-import torch.autograd as autograd
-import torch.cuda.comm as comm
-from torch.autograd.function import once_differentiable
-from torch.utils.cpp_extension import load
-
-_src_path = path.join(path.dirname(path.abspath(__file__)), "src")
-_backend = load(name="inplace_abn",
-                extra_cflags=["-O3"],
-                sources=[path.join(_src_path, f) for f in [
-                    "inplace_abn.cpp",
-                    "inplace_abn_cpu.cpp",
-                    "inplace_abn_cuda.cu",
-                    "inplace_abn_cuda_half.cu"
-                ]],
-                extra_cuda_cflags=["--expt-extended-lambda"])
-
-# Activation names
-ACT_RELU = "relu"
-ACT_LEAKY_RELU = "leaky_relu"
-ACT_ELU = "elu"
-ACT_NONE = "none"
-
-
-def _check(fn, *args, **kwargs):
-    success = fn(*args, **kwargs)
-    if not success:
-        raise RuntimeError("CUDA Error encountered in {}".format(fn))
-
-
-def _broadcast_shape(x):
-    out_size = []
-    for i, s in enumerate(x.size()):
-        if i != 1:
-            out_size.append(1)
-        else:
-            out_size.append(s)
-    return out_size
-
-
-def _reduce(x):
-    if len(x.size()) == 2:
-        return x.sum(dim=0)
-    else:
-        n, c = x.size()[0:2]
-        return x.contiguous().view((n, c, -1)).sum(2).sum(0)
-
-
-def _count_samples(x):
-    count = 1
-    for i, s in enumerate(x.size()):
-        if i != 1:
-            count *= s
-    return count
-
-
-def _act_forward(ctx, x):
-    if ctx.activation == ACT_LEAKY_RELU:
-        _backend.leaky_relu_forward(x, ctx.slope)
-    elif ctx.activation == ACT_ELU:
-        _backend.elu_forward(x)
-    elif ctx.activation == ACT_NONE:
-        pass
-
-
-def _act_backward(ctx, x, dx):
-    if ctx.activation == ACT_LEAKY_RELU:
-        _backend.leaky_relu_backward(x, dx, ctx.slope)
-    elif ctx.activation == ACT_ELU:
-        _backend.elu_backward(x, dx)
-    elif ctx.activation == ACT_NONE:
-        pass
-
-
-class InPlaceABN(autograd.Function):
-    @staticmethod
-    def forward(ctx, x, weight, bias, running_mean, running_var,
-                training=True, momentum=0.1, eps=1e-05, activation=ACT_LEAKY_RELU, slope=0.01):
-        # Save context
-        ctx.training = training
-        ctx.momentum = momentum
-        ctx.eps = eps
-        ctx.activation = activation
-        ctx.slope = slope
-        ctx.affine = weight is not None and bias is not None
-
-        # Prepare inputs
-        count = _count_samples(x)
-        x = x.contiguous()
-        weight = weight.contiguous() if ctx.affine else x.new_empty(0)
-        bias = bias.contiguous() if ctx.affine else x.new_empty(0)
-
-        if ctx.training:
-            mean, var = _backend.mean_var(x)
-
-            # Update running stats
-            running_mean.mul_((1 - ctx.momentum)).add_(ctx.momentum * mean)
-            running_var.mul_((1 - ctx.momentum)).add_(ctx.momentum * var * count / (count - 1))
-
-            # Mark in-place modified tensors
-            ctx.mark_dirty(x, running_mean, running_var)
-        else:
-            mean, var = running_mean.contiguous(), running_var.contiguous()
-            ctx.mark_dirty(x)
-
-        # BN forward + activation
-        _backend.forward(x, mean, var, weight, bias, ctx.affine, ctx.eps)
-        _act_forward(ctx, x)
-
-        # Output
-        ctx.var = var
-        ctx.save_for_backward(x, var, weight, bias)
-        ctx.mark_non_differentiable(running_mean, running_var)
-        return x, running_mean, running_var
-
-    @staticmethod
-    @once_differentiable
-    def backward(ctx, dz, _drunning_mean, _drunning_var):
-        z, var, weight, bias = ctx.saved_tensors
-        dz = dz.contiguous()
-
-        # Undo activation
-        _act_backward(ctx, z, dz)
-
-        if ctx.training:
-            edz, eydz = _backend.edz_eydz(z, dz, weight, bias, ctx.affine, ctx.eps)
-        else:
-            # TODO: implement simplified CUDA backward for inference mode
-            edz = dz.new_zeros(dz.size(1))
-            eydz = dz.new_zeros(dz.size(1))
-
-        dx = _backend.backward(z, dz, var, weight, bias, edz, eydz, ctx.affine, ctx.eps)
-        # dweight = eydz * weight.sign() if ctx.affine else None
-        dweight = eydz if ctx.affine else None
-        if dweight is not None:
-            dweight[weight < 0] *= -1
-        dbias = edz if ctx.affine else None
-
-        return dx, dweight, dbias, None, None, None, None, None, None, None
-
-
-class InPlaceABNSync(autograd.Function):
-    @classmethod
-    def forward(cls, ctx, x, weight, bias, running_mean, running_var,
-                training=True, momentum=0.1, eps=1e-05, activation=ACT_LEAKY_RELU, slope=0.01, equal_batches=True):
-        # Save context
-        ctx.training = training
-        ctx.momentum = momentum
-        ctx.eps = eps
-        ctx.activation = activation
-        ctx.slope = slope
-        ctx.affine = weight is not None and bias is not None
-
-        # Prepare inputs
-        ctx.world_size = dist.get_world_size() if dist.is_initialized() else 1
-
-        # count = _count_samples(x)
-        batch_size = x.new_tensor([x.shape[0]], dtype=torch.long)
-
-        x = x.contiguous()
-        weight = weight.contiguous() if ctx.affine else x.new_empty(0)
-        bias = bias.contiguous() if ctx.affine else x.new_empty(0)
-
-        if ctx.training:
-            mean, var = _backend.mean_var(x)
-            if ctx.world_size > 1:
-                # get global batch size
-                if equal_batches:
-                    batch_size *= ctx.world_size
-                else:
-                    dist.all_reduce(batch_size, dist.ReduceOp.SUM)
-
-                ctx.factor = x.shape[0] / float(batch_size.item())
-
-                mean_all = mean.clone() * ctx.factor
-                dist.all_reduce(mean_all, dist.ReduceOp.SUM)
-
-                var_all = (var + (mean - mean_all) ** 2) * ctx.factor
-                dist.all_reduce(var_all, dist.ReduceOp.SUM)
-
-                mean = mean_all
-                var = var_all
-
-            # Update running stats
-            running_mean.mul_((1 - ctx.momentum)).add_(ctx.momentum * mean)
-            count = batch_size.item() * x.view(x.shape[0], x.shape[1], -1).shape[-1]
-            running_var.mul_((1 - ctx.momentum)).add_(ctx.momentum * var * (float(count) / (count - 1)))
-
-            # Mark in-place modified tensors
-            ctx.mark_dirty(x, running_mean, running_var)
-        else:
-            mean, var = running_mean.contiguous(), running_var.contiguous()
-            ctx.mark_dirty(x)
-
-        # BN forward + activation
-        _backend.forward(x, mean, var, weight, bias, ctx.affine, ctx.eps)
-        _act_forward(ctx, x)
-
-        # Output
-        ctx.var = var
-        ctx.save_for_backward(x, var, weight, bias)
-        ctx.mark_non_differentiable(running_mean, running_var)
-        return x, running_mean, running_var
-
-    @staticmethod
-    @once_differentiable
-    def backward(ctx, dz, _drunning_mean, _drunning_var):
-        z, var, weight, bias = ctx.saved_tensors
-        dz = dz.contiguous()
-
-        # Undo activation
-        _act_backward(ctx, z, dz)
-
-        if ctx.training:
-            edz, eydz = _backend.edz_eydz(z, dz, weight, bias, ctx.affine, ctx.eps)
-            edz_local = edz.clone()
-            eydz_local = eydz.clone()
-
-            if ctx.world_size > 1:
-                edz *= ctx.factor
-                dist.all_reduce(edz, dist.ReduceOp.SUM)
-
-                eydz *= ctx.factor
-                dist.all_reduce(eydz, dist.ReduceOp.SUM)
-        else:
-            edz_local = edz = dz.new_zeros(dz.size(1))
-            eydz_local = eydz = dz.new_zeros(dz.size(1))
-
-        dx = _backend.backward(z, dz, var, weight, bias, edz, eydz, ctx.affine, ctx.eps)
-        # dweight = eydz_local * weight.sign() if ctx.affine else None
-        dweight = eydz_local if ctx.affine else None
-        if dweight is not None:
-            dweight[weight < 0] *= -1
-        dbias = edz_local if ctx.affine else None
-
-        return dx, dweight, dbias, None, None, None, None, None, None, None
-
-
-inplace_abn = InPlaceABN.apply
-inplace_abn_sync = InPlaceABNSync.apply
-
-__all__ = ["inplace_abn", "inplace_abn_sync", "ACT_RELU", "ACT_LEAKY_RELU", "ACT_ELU", "ACT_NONE"]

idm_masking/humanparsing/modules/misc.py

@@ -1,21 +0,0 @@
-import torch.nn as nn
-import torch
-import torch.distributed as dist
-
-class GlobalAvgPool2d(nn.Module):
-    def __init__(self):
-        """Global average pooling over the input's spatial dimensions"""
-        super(GlobalAvgPool2d, self).__init__()
-
-    def forward(self, inputs):
-        in_size = inputs.size()
-        return inputs.view((in_size[0], in_size[1], -1)).mean(dim=2)
-
-class SingleGPU(nn.Module):
-    def __init__(self, module):
-        super(SingleGPU, self).__init__()
-        self.module=module
-
-    def forward(self, input):
-        return self.module(input.cuda(non_blocking=True))
-

idm_masking/humanparsing/modules/residual.py

@@ -1,182 +0,0 @@
-from collections import OrderedDict
-
-import torch.nn as nn
-
-from .bn import ABN, ACT_LEAKY_RELU, ACT_ELU, ACT_NONE
-import torch.nn.functional as functional
-
-
-class ResidualBlock(nn.Module):
-    """Configurable residual block
-
-    Parameters
-    ----------
-    in_channels : int
-        Number of input channels.
-    channels : list of int
-        Number of channels in the internal feature maps. Can either have two or three elements: if three construct
-        a residual block with two `3 x 3` convolutions, otherwise construct a bottleneck block with `1 x 1`, then
-        `3 x 3` then `1 x 1` convolutions.
-    stride : int
-        Stride of the first `3 x 3` convolution
-    dilation : int
-        Dilation to apply to the `3 x 3` convolutions.
-    groups : int
-        Number of convolution groups. This is used to create ResNeXt-style blocks and is only compatible with
-        bottleneck blocks.
-    norm_act : callable
-        Function to create normalization / activation Module.
-    dropout: callable
-        Function to create Dropout Module.
-    """
-
-    def __init__(self,
-                 in_channels,
-                 channels,
-                 stride=1,
-                 dilation=1,
-                 groups=1,
-                 norm_act=ABN,
-                 dropout=None):
-        super(ResidualBlock, self).__init__()
-
-        # Check parameters for inconsistencies
-        if len(channels) != 2 and len(channels) != 3:
-            raise ValueError("channels must contain either two or three values")
-        if len(channels) == 2 and groups != 1:
-            raise ValueError("groups > 1 are only valid if len(channels) == 3")
-
-        is_bottleneck = len(channels) == 3
-        need_proj_conv = stride != 1 or in_channels != channels[-1]
-
-        if not is_bottleneck:
-            bn2 = norm_act(channels[1])
-            bn2.activation = ACT_NONE
-            layers = [
-                ("conv1", nn.Conv2d(in_channels, channels[0], 3, stride=stride, padding=dilation, bias=False,
-                                    dilation=dilation)),
-                ("bn1", norm_act(channels[0])),
-                ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=1, padding=dilation, bias=False,
-                                    dilation=dilation)),
-                ("bn2", bn2)
-            ]
-            if dropout is not None:
-                layers = layers[0:2] + [("dropout", dropout())] + layers[2:]
-        else:
-            bn3 = norm_act(channels[2])
-            bn3.activation = ACT_NONE
-            layers = [
-                ("conv1", nn.Conv2d(in_channels, channels[0], 1, stride=1, padding=0, bias=False)),
-                ("bn1", norm_act(channels[0])),
-                ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=stride, padding=dilation, bias=False,
-                                    groups=groups, dilation=dilation)),
-                ("bn2", norm_act(channels[1])),
-                ("conv3", nn.Conv2d(channels[1], channels[2], 1, stride=1, padding=0, bias=False)),
-                ("bn3", bn3)
-            ]
-            if dropout is not None:
-                layers = layers[0:4] + [("dropout", dropout())] + layers[4:]
-        self.convs = nn.Sequential(OrderedDict(layers))
-
-        if need_proj_conv:
-            self.proj_conv = nn.Conv2d(in_channels, channels[-1], 1, stride=stride, padding=0, bias=False)
-            self.proj_bn = norm_act(channels[-1])
-            self.proj_bn.activation = ACT_NONE
-
-    def forward(self, x):
-        if hasattr(self, "proj_conv"):
-            residual = self.proj_conv(x)
-            residual = self.proj_bn(residual)
-        else:
-            residual = x
-        x = self.convs(x) + residual
-
-        if self.convs.bn1.activation == ACT_LEAKY_RELU:
-            return functional.leaky_relu(x, negative_slope=self.convs.bn1.slope, inplace=True)
-        elif self.convs.bn1.activation == ACT_ELU:
-            return functional.elu(x, inplace=True)
-        else:
-            return x
-
-
-class IdentityResidualBlock(nn.Module):
-    def __init__(self,
-                 in_channels,
-                 channels,
-                 stride=1,
-                 dilation=1,
-                 groups=1,
-                 norm_act=ABN,
-                 dropout=None):
-        """Configurable identity-mapping residual block
-
-        Parameters
-        ----------
-        in_channels : int
-            Number of input channels.
-        channels : list of int
-            Number of channels in the internal feature maps. Can either have two or three elements: if three construct
-            a residual block with two `3 x 3` convolutions, otherwise construct a bottleneck block with `1 x 1`, then
-            `3 x 3` then `1 x 1` convolutions.
-        stride : int
-            Stride of the first `3 x 3` convolution
-        dilation : int
-            Dilation to apply to the `3 x 3` convolutions.
-        groups : int
-            Number of convolution groups. This is used to create ResNeXt-style blocks and is only compatible with
-            bottleneck blocks.
-        norm_act : callable
-            Function to create normalization / activation Module.
-        dropout: callable
-            Function to create Dropout Module.
-        """
-        super(IdentityResidualBlock, self).__init__()
-
-        # Check parameters for inconsistencies
-        if len(channels) != 2 and len(channels) != 3:
-            raise ValueError("channels must contain either two or three values")
-        if len(channels) == 2 and groups != 1:
-            raise ValueError("groups > 1 are only valid if len(channels) == 3")
-
-        is_bottleneck = len(channels) == 3
-        need_proj_conv = stride != 1 or in_channels != channels[-1]
-
-        self.bn1 = norm_act(in_channels)
-        if not is_bottleneck:
-            layers = [
-                ("conv1", nn.Conv2d(in_channels, channels[0], 3, stride=stride, padding=dilation, bias=False,
-                                    dilation=dilation)),
-                ("bn2", norm_act(channels[0])),
-                ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=1, padding=dilation, bias=False,
-                                    dilation=dilation))
-            ]
-            if dropout is not None:
-                layers = layers[0:2] + [("dropout", dropout())] + layers[2:]
-        else:
-            layers = [
-                ("conv1", nn.Conv2d(in_channels, channels[0], 1, stride=stride, padding=0, bias=False)),
-                ("bn2", norm_act(channels[0])),
-                ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=1, padding=dilation, bias=False,
-                                    groups=groups, dilation=dilation)),
-                ("bn3", norm_act(channels[1])),
-                ("conv3", nn.Conv2d(channels[1], channels[2], 1, stride=1, padding=0, bias=False))
-            ]
-            if dropout is not None:
-                layers = layers[0:4] + [("dropout", dropout())] + layers[4:]
-        self.convs = nn.Sequential(OrderedDict(layers))
-
-        if need_proj_conv:
-            self.proj_conv = nn.Conv2d(in_channels, channels[-1], 1, stride=stride, padding=0, bias=False)
-
-    def forward(self, x):
-        if hasattr(self, "proj_conv"):
-            bn1 = self.bn1(x)
-            shortcut = self.proj_conv(bn1)
-        else:
-            shortcut = x.clone()
-            bn1 = self.bn1(x)
-
-        out = self.convs(bn1)
-        out.add_(shortcut)
-
-        return out

idm_masking/humanparsing/modules/src/checks.h

@@ -1,15 +0,0 @@
-#pragma once
-
-#include <ATen/ATen.h>
-
-// Define AT_CHECK for old version of ATen where the same function was called AT_ASSERT
-#ifndef AT_CHECK
-#define AT_CHECK AT_ASSERT
-#endif
-
-#define CHECK_CUDA(x) AT_CHECK((x).type().is_cuda(), #x " must be a CUDA tensor")
-#define CHECK_CPU(x) AT_CHECK(!(x).type().is_cuda(), #x " must be a CPU tensor")
-#define CHECK_CONTIGUOUS(x) AT_CHECK((x).is_contiguous(), #x " must be contiguous")
-
-#define CHECK_CUDA_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
-#define CHECK_CPU_INPUT(x) CHECK_CPU(x); CHECK_CONTIGUOUS(x)

idm_masking/humanparsing/modules/src/inplace_abn.cpp

@@ -1,95 +0,0 @@
-#include <torch/extension.h>
-
-#include <vector>
-
-#include "inplace_abn.h"
-
-std::vector<at::Tensor> mean_var(at::Tensor x) {
-  if (x.is_cuda()) {
-    if (x.type().scalarType() == at::ScalarType::Half) {
-      return mean_var_cuda_h(x);
-    } else {
-      return mean_var_cuda(x);
-    }
-  } else {
-    return mean_var_cpu(x);
-  }
-}
-
-at::Tensor forward(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias,
-                   bool affine, float eps) {
-  if (x.is_cuda()) {
-    if (x.type().scalarType() == at::ScalarType::Half) {
-      return forward_cuda_h(x, mean, var, weight, bias, affine, eps);
-    } else {
-      return forward_cuda(x, mean, var, weight, bias, affine, eps);
-    }
-  } else {
-    return forward_cpu(x, mean, var, weight, bias, affine, eps);
-  }
-}
-
-std::vector<at::Tensor> edz_eydz(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias,
-                                 bool affine, float eps) {
-  if (z.is_cuda()) {
-    if (z.type().scalarType() == at::ScalarType::Half) {
-      return edz_eydz_cuda_h(z, dz, weight, bias, affine, eps);
-    } else {
-      return edz_eydz_cuda(z, dz, weight, bias, affine, eps);
-	}
-  } else {
-    return edz_eydz_cpu(z, dz, weight, bias, affine, eps);
-  }
-}
-
-at::Tensor backward(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias,
-                                 at::Tensor edz, at::Tensor eydz, bool affine, float eps) {
-  if (z.is_cuda()) {
-    if (z.type().scalarType() == at::ScalarType::Half) {
-      return backward_cuda_h(z, dz, var, weight, bias, edz, eydz, affine, eps);
-	} else {
-      return backward_cuda(z, dz, var, weight, bias, edz, eydz, affine, eps);
-    }
-  } else {
-    return backward_cpu(z, dz, var, weight, bias, edz, eydz, affine, eps);
-  }
-}
-
-void leaky_relu_forward(at::Tensor z, float slope) {
-  at::leaky_relu_(z, slope);
-}
-
-void leaky_relu_backward(at::Tensor z, at::Tensor dz, float slope) {
-  if (z.is_cuda()) {
-    if (z.type().scalarType() == at::ScalarType::Half) {
-      return leaky_relu_backward_cuda_h(z, dz, slope);
-	} else {
-      return leaky_relu_backward_cuda(z, dz, slope);
-    }
-  } else {
-    return leaky_relu_backward_cpu(z, dz, slope);
-  }
-}
-
-void elu_forward(at::Tensor z) {
-  at::elu_(z);
-}
-
-void elu_backward(at::Tensor z, at::Tensor dz) {
-  if (z.is_cuda()) {
-    return elu_backward_cuda(z, dz);
-  } else {
-    return elu_backward_cpu(z, dz);
-  }
-}
-
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("mean_var", &mean_var, "Mean and variance computation");
-  m.def("forward", &forward, "In-place forward computation");
-  m.def("edz_eydz", &edz_eydz, "First part of backward computation");
-  m.def("backward", &backward, "Second part of backward computation");
-  m.def("leaky_relu_forward", &leaky_relu_forward, "Leaky relu forward computation");
-  m.def("leaky_relu_backward", &leaky_relu_backward, "Leaky relu backward computation and inversion");
-  m.def("elu_forward", &elu_forward, "Elu forward computation");
-  m.def("elu_backward", &elu_backward, "Elu backward computation and inversion");
-}

idm_masking/humanparsing/modules/src/inplace_abn.h

@@ -1,88 +0,0 @@
-#pragma once
-
-#include <ATen/ATen.h>
-
-#include <vector>
-
-std::vector<at::Tensor> mean_var_cpu(at::Tensor x);
-std::vector<at::Tensor> mean_var_cuda(at::Tensor x);
-std::vector<at::Tensor> mean_var_cuda_h(at::Tensor x);
-
-at::Tensor forward_cpu(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias,
-                       bool affine, float eps);
-at::Tensor forward_cuda(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias,
-                        bool affine, float eps);
-at::Tensor forward_cuda_h(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias,
-                          bool affine, float eps);
-
-std::vector<at::Tensor> edz_eydz_cpu(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias,
-                                     bool affine, float eps);
-std::vector<at::Tensor> edz_eydz_cuda(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias,
-                                      bool affine, float eps);
-std::vector<at::Tensor> edz_eydz_cuda_h(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias,
-                                        bool affine, float eps);
-
-at::Tensor backward_cpu(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias,
-                                     at::Tensor edz, at::Tensor eydz, bool affine, float eps);
-at::Tensor backward_cuda(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias,
-                                      at::Tensor edz, at::Tensor eydz, bool affine, float eps);
-at::Tensor backward_cuda_h(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias,
-                                        at::Tensor edz, at::Tensor eydz, bool affine, float eps);
-
-void leaky_relu_backward_cpu(at::Tensor z, at::Tensor dz, float slope);
-void leaky_relu_backward_cuda(at::Tensor z, at::Tensor dz, float slope);
-void leaky_relu_backward_cuda_h(at::Tensor z, at::Tensor dz, float slope);
-
-void elu_backward_cpu(at::Tensor z, at::Tensor dz);
-void elu_backward_cuda(at::Tensor z, at::Tensor dz);
-
-static void get_dims(at::Tensor x, int64_t& num, int64_t& chn, int64_t& sp) {
-  num = x.size(0);
-  chn = x.size(1);
-  sp = 1;
-  for (int64_t i = 2; i < x.ndimension(); ++i)
-    sp *= x.size(i);
-}
-
-/*
- * Specialized CUDA reduction functions for BN
- */
-#ifdef __CUDACC__
-
-#include "utils/cuda.cuh"
-
-template <typename T, typename Op>
-__device__ T reduce(Op op, int plane, int N, int S) {
-  T sum = (T)0;
-  for (int batch = 0; batch < N; ++batch) {
-    for (int x = threadIdx.x; x < S; x += blockDim.x) {
-      sum += op(batch, plane, x);
-    }
-  }
-
-  // sum over NumThreads within a warp
-  sum = warpSum(sum);
-
-  // 'transpose', and reduce within warp again
-  __shared__ T shared[32];
-  __syncthreads();
-  if (threadIdx.x % WARP_SIZE == 0) {
-    shared[threadIdx.x / WARP_SIZE] = sum;
-  }
-  if (threadIdx.x >= blockDim.x / WARP_SIZE && threadIdx.x < WARP_SIZE) {
-    // zero out the other entries in shared
-    shared[threadIdx.x] = (T)0;
-  }
-  __syncthreads();
-  if (threadIdx.x / WARP_SIZE == 0) {
-    sum = warpSum(shared[threadIdx.x]);
-    if (threadIdx.x == 0) {
-      shared[0] = sum;
-    }
-  }
-  __syncthreads();
-
-  // Everyone picks it up, should be broadcast into the whole gradInput
-  return shared[0];
-}
-#endif

idm_masking/humanparsing/modules/src/inplace_abn_cpu.cpp

@@ -1,119 +0,0 @@
-#include <ATen/ATen.h>
-
-#include <vector>
-
-#include "utils/checks.h"
-#include "inplace_abn.h"
-
-at::Tensor reduce_sum(at::Tensor x) {
-  if (x.ndimension() == 2) {
-    return x.sum(0);
-  } else {
-    auto x_view = x.view({x.size(0), x.size(1), -1});
-    return x_view.sum(-1).sum(0);
-  }
-}
-
-at::Tensor broadcast_to(at::Tensor v, at::Tensor x) {
-  if (x.ndimension() == 2) {
-    return v;
-  } else {
-    std::vector<int64_t> broadcast_size = {1, -1};
-    for (int64_t i = 2; i < x.ndimension(); ++i)
-      broadcast_size.push_back(1);
-
-    return v.view(broadcast_size);
-  }
-}
-
-int64_t count(at::Tensor x) {
-  int64_t count = x.size(0);
-  for (int64_t i = 2; i < x.ndimension(); ++i)
-    count *= x.size(i);
-
-  return count;
-}
-
-at::Tensor invert_affine(at::Tensor z, at::Tensor weight, at::Tensor bias, bool affine, float eps) {
-  if (affine) {
-    return (z - broadcast_to(bias, z)) / broadcast_to(at::abs(weight) + eps, z);
-  } else {
-    return z;
-  }
-}
-
-std::vector<at::Tensor> mean_var_cpu(at::Tensor x) {
-  auto num = count(x);
-  auto mean = reduce_sum(x) / num;
-  auto diff = x - broadcast_to(mean, x);
-  auto var = reduce_sum(diff.pow(2)) / num;
-
-  return {mean, var};
-}
-
-at::Tensor forward_cpu(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias,
-                       bool affine, float eps) {
-  auto gamma = affine ? at::abs(weight) + eps : at::ones_like(var);
-  auto mul = at::rsqrt(var + eps) * gamma;
-
-  x.sub_(broadcast_to(mean, x));
-  x.mul_(broadcast_to(mul, x));
-  if (affine) x.add_(broadcast_to(bias, x));
-
-  return x;
-}
-
-std::vector<at::Tensor> edz_eydz_cpu(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias,
-                                     bool affine, float eps) {
-  auto edz = reduce_sum(dz);
-  auto y = invert_affine(z, weight, bias, affine, eps);
-  auto eydz = reduce_sum(y * dz);
-
-  return {edz, eydz};
-}
-
-at::Tensor backward_cpu(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias,
-                                     at::Tensor edz, at::Tensor eydz, bool affine, float eps) {
-  auto y = invert_affine(z, weight, bias, affine, eps);
-  auto mul = affine ? at::rsqrt(var + eps) * (at::abs(weight) + eps) : at::rsqrt(var + eps);
-
-  auto num = count(z);
-  auto dx = (dz - broadcast_to(edz / num, dz) - y * broadcast_to(eydz / num, dz)) * broadcast_to(mul, dz);
-  return dx;
-}
-
-void leaky_relu_backward_cpu(at::Tensor z, at::Tensor dz, float slope) {
-  CHECK_CPU_INPUT(z);
-  CHECK_CPU_INPUT(dz);
-
-  AT_DISPATCH_FLOATING_TYPES(z.type(), "leaky_relu_backward_cpu", ([&] {
-    int64_t count = z.numel();
-    auto *_z = z.data<scalar_t>();
-    auto *_dz = dz.data<scalar_t>();
-
-    for (int64_t i = 0; i < count; ++i) {
-      if (_z[i] < 0) {
-        _z[i] *= 1 / slope;
-        _dz[i] *= slope;
-      }
-    }
-  }));
-}
-
-void elu_backward_cpu(at::Tensor z, at::Tensor dz) {
-  CHECK_CPU_INPUT(z);
-  CHECK_CPU_INPUT(dz);
-
-  AT_DISPATCH_FLOATING_TYPES(z.type(), "elu_backward_cpu", ([&] {
-    int64_t count = z.numel();
-    auto *_z = z.data<scalar_t>();
-    auto *_dz = dz.data<scalar_t>();
-
-    for (int64_t i = 0; i < count; ++i) {
-      if (_z[i] < 0) {
-        _z[i] = log1p(_z[i]);
-        _dz[i] *= (_z[i] + 1.f);
-      }
-    }
-  }));
-}

idm_masking/humanparsing/modules/src/inplace_abn_cuda.cu

@@ -1,333 +0,0 @@
-#include <ATen/ATen.h>
-
-#include <thrust/device_ptr.h>
-#include <thrust/transform.h>
-
-#include <vector>
-
-#include "utils/checks.h"
-#include "utils/cuda.cuh"
-#include "inplace_abn.h"
-
-#include <ATen/cuda/CUDAContext.h>
-
-// Operations for reduce
-template<typename T>
-struct SumOp {
-  __device__ SumOp(const T *t, int c, int s)
-      : tensor(t), chn(c), sp(s) {}
-  __device__ __forceinline__ T operator()(int batch, int plane, int n) {
-    return tensor[(batch * chn + plane) * sp + n];
-  }
-  const T *tensor;
-  const int chn;
-  const int sp;
-};
-
-template<typename T>
-struct VarOp {
-  __device__ VarOp(T m, const T *t, int c, int s)
-      : mean(m), tensor(t), chn(c), sp(s) {}
-  __device__ __forceinline__ T operator()(int batch, int plane, int n) {
-    T val = tensor[(batch * chn + plane) * sp + n];
-    return (val - mean) * (val - mean);
-  }
-  const T mean;
-  const T *tensor;
-  const int chn;
-  const int sp;
-};
-
-template<typename T>
-struct GradOp {
-  __device__ GradOp(T _weight, T _bias, const T *_z, const T *_dz, int c, int s)
-      : weight(_weight), bias(_bias), z(_z), dz(_dz), chn(c), sp(s) {}
-  __device__ __forceinline__ Pair<T> operator()(int batch, int plane, int n) {
-    T _y = (z[(batch * chn + plane) * sp + n] - bias) / weight;
-    T _dz = dz[(batch * chn + plane) * sp + n];
-    return Pair<T>(_dz, _y * _dz);
-  }
-  const T weight;
-  const T bias;
-  const T *z;
-  const T *dz;
-  const int chn;
-  const int sp;
-};
-
-/***********
- * mean_var
- ***********/
-
-template<typename T>
-__global__ void mean_var_kernel(const T *x, T *mean, T *var, int num, int chn, int sp) {
-  int plane = blockIdx.x;
-  T norm = T(1) / T(num * sp);
-
-  T _mean = reduce<T, SumOp<T>>(SumOp<T>(x, chn, sp), plane, num, sp) * norm;
-  __syncthreads();
-  T _var = reduce<T, VarOp<T>>(VarOp<T>(_mean, x, chn, sp), plane, num, sp) * norm;
-
-  if (threadIdx.x == 0) {
-    mean[plane] = _mean;
-    var[plane] = _var;
-  }
-}
-
-std::vector<at::Tensor> mean_var_cuda(at::Tensor x) {
-  CHECK_CUDA_INPUT(x);
-
-  // Extract dimensions
-  int64_t num, chn, sp;
-  get_dims(x, num, chn, sp);
-
-  // Prepare output tensors
-  auto mean = at::empty({chn}, x.options());
-  auto var = at::empty({chn}, x.options());
-
-  // Run kernel
-  dim3 blocks(chn);
-  dim3 threads(getNumThreads(sp));
-  auto stream = at::cuda::getCurrentCUDAStream();
-  AT_DISPATCH_FLOATING_TYPES(x.type(), "mean_var_cuda", ([&] {
-    mean_var_kernel<scalar_t><<<blocks, threads, 0, stream>>>(
-        x.data<scalar_t>(),
-        mean.data<scalar_t>(),
-        var.data<scalar_t>(),
-        num, chn, sp);
-  }));
-
-  return {mean, var};
-}
-
-/**********
- * forward
- **********/
-
-template<typename T>
-__global__ void forward_kernel(T *x, const T *mean, const T *var, const T *weight, const T *bias,
-                               bool affine, float eps, int num, int chn, int sp) {
-  int plane = blockIdx.x;
-
-  T _mean = mean[plane];
-  T _var = var[plane];
-  T _weight = affine ? abs(weight[plane]) + eps : T(1);
-  T _bias = affine ? bias[plane] : T(0);
-
-  T mul = rsqrt(_var + eps) * _weight;
-
-  for (int batch = 0; batch < num; ++batch) {
-    for (int n = threadIdx.x; n < sp; n += blockDim.x) {
-      T _x = x[(batch * chn + plane) * sp + n];
-      T _y = (_x - _mean) * mul + _bias;
-
-      x[(batch * chn + plane) * sp + n] = _y;
-    }
-  }
-}
-
-at::Tensor forward_cuda(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias,
-                        bool affine, float eps) {
-  CHECK_CUDA_INPUT(x);
-  CHECK_CUDA_INPUT(mean);
-  CHECK_CUDA_INPUT(var);
-  CHECK_CUDA_INPUT(weight);
-  CHECK_CUDA_INPUT(bias);
-
-  // Extract dimensions
-  int64_t num, chn, sp;
-  get_dims(x, num, chn, sp);
-
-  // Run kernel
-  dim3 blocks(chn);
-  dim3 threads(getNumThreads(sp));
-  auto stream = at::cuda::getCurrentCUDAStream();
-  AT_DISPATCH_FLOATING_TYPES(x.type(), "forward_cuda", ([&] {
-    forward_kernel<scalar_t><<<blocks, threads, 0, stream>>>(
-        x.data<scalar_t>(),
-        mean.data<scalar_t>(),
-        var.data<scalar_t>(),
-        weight.data<scalar_t>(),
-        bias.data<scalar_t>(),
-        affine, eps, num, chn, sp);
-  }));
-
-  return x;
-}
-
-/***********
- * edz_eydz
- ***********/
-
-template<typename T>
-__global__ void edz_eydz_kernel(const T *z, const T *dz, const T *weight, const T *bias,
-                                T *edz, T *eydz, bool affine, float eps, int num, int chn, int sp) {
-  int plane = blockIdx.x;
-
-  T _weight = affine ? abs(weight[plane]) + eps : 1.f;
-  T _bias = affine ? bias[plane] : 0.f;
-
-  Pair<T> res = reduce<Pair<T>, GradOp<T>>(GradOp<T>(_weight, _bias, z, dz, chn, sp), plane, num, sp);
-  __syncthreads();
-
-  if (threadIdx.x == 0) {
-    edz[plane] = res.v1;
-    eydz[plane] = res.v2;
-  }
-}
-
-std::vector<at::Tensor> edz_eydz_cuda(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias,
-                                      bool affine, float eps) {
-  CHECK_CUDA_INPUT(z);
-  CHECK_CUDA_INPUT(dz);
-  CHECK_CUDA_INPUT(weight);
-  CHECK_CUDA_INPUT(bias);
-
-  // Extract dimensions
-  int64_t num, chn, sp;
-  get_dims(z, num, chn, sp);
-
-  auto edz = at::empty({chn}, z.options());
-  auto eydz = at::empty({chn}, z.options());
-
-  // Run kernel
-  dim3 blocks(chn);
-  dim3 threads(getNumThreads(sp));
-  auto stream = at::cuda::getCurrentCUDAStream();
-  AT_DISPATCH_FLOATING_TYPES(z.type(), "edz_eydz_cuda", ([&] {
-    edz_eydz_kernel<scalar_t><<<blocks, threads, 0, stream>>>(
-        z.data<scalar_t>(),
-        dz.data<scalar_t>(),
-        weight.data<scalar_t>(),
-        bias.data<scalar_t>(),
-        edz.data<scalar_t>(),
-        eydz.data<scalar_t>(),
-        affine, eps, num, chn, sp);
-  }));
-
-  return {edz, eydz};
-}
-
-/***********
- * backward
- ***********/
-
-template<typename T>
-__global__ void backward_kernel(const T *z, const T *dz, const T *var, const T *weight, const T *bias, const T *edz,
-	                        const T *eydz, T *dx, bool affine, float eps, int num, int chn, int sp) {
-  int plane = blockIdx.x;
-
-  T _weight = affine ? abs(weight[plane]) + eps : 1.f;
-  T _bias = affine ? bias[plane] : 0.f;
-  T _var = var[plane];
-  T _edz = edz[plane];
-  T _eydz = eydz[plane];
-
-  T _mul = _weight * rsqrt(_var + eps);
-  T count = T(num * sp);
-
-  for (int batch = 0; batch < num; ++batch) {
-    for (int n = threadIdx.x; n < sp; n += blockDim.x) {
-      T _dz = dz[(batch * chn + plane) * sp + n];
-      T _y = (z[(batch * chn + plane) * sp + n] - _bias) / _weight;
-
-      dx[(batch * chn + plane) * sp + n] = (_dz - _edz / count - _y * _eydz / count) * _mul;
-    }
-  }
-}
-
-at::Tensor backward_cuda(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias,
-                                      at::Tensor edz, at::Tensor eydz, bool affine, float eps) {
-  CHECK_CUDA_INPUT(z);
-  CHECK_CUDA_INPUT(dz);
-  CHECK_CUDA_INPUT(var);
-  CHECK_CUDA_INPUT(weight);
-  CHECK_CUDA_INPUT(bias);
-  CHECK_CUDA_INPUT(edz);
-  CHECK_CUDA_INPUT(eydz);
-
-  // Extract dimensions
-  int64_t num, chn, sp;
-  get_dims(z, num, chn, sp);
-
-  auto dx = at::zeros_like(z);
-
-  // Run kernel
-  dim3 blocks(chn);
-  dim3 threads(getNumThreads(sp));
-  auto stream = at::cuda::getCurrentCUDAStream();
-  AT_DISPATCH_FLOATING_TYPES(z.type(), "backward_cuda", ([&] {
-    backward_kernel<scalar_t><<<blocks, threads, 0, stream>>>(
-        z.data<scalar_t>(),
-        dz.data<scalar_t>(),
-        var.data<scalar_t>(),
-        weight.data<scalar_t>(),
-        bias.data<scalar_t>(),
-        edz.data<scalar_t>(),
-        eydz.data<scalar_t>(),
-        dx.data<scalar_t>(),
-        affine, eps, num, chn, sp);
-  }));
-
-  return dx;
-}
-
-/**************
- * activations
- **************/
-
-template<typename T>
-inline void leaky_relu_backward_impl(T *z, T *dz, float slope, int64_t count) {
-  // Create thrust pointers
-  thrust::device_ptr<T> th_z = thrust::device_pointer_cast(z);
-  thrust::device_ptr<T> th_dz = thrust::device_pointer_cast(dz);
-
-  auto stream = at::cuda::getCurrentCUDAStream();
-  thrust::transform_if(thrust::cuda::par.on(stream),
-                       th_dz, th_dz + count, th_z, th_dz,
-                       [slope] __device__ (const T& dz) { return dz * slope; },
-                       [] __device__ (const T& z) { return z < 0; });
-  thrust::transform_if(thrust::cuda::par.on(stream),
-                       th_z, th_z + count, th_z,
-                       [slope] __device__ (const T& z) { return z / slope; },
-                       [] __device__ (const T& z) { return z < 0; });
-}
-
-void leaky_relu_backward_cuda(at::Tensor z, at::Tensor dz, float slope) {
-  CHECK_CUDA_INPUT(z);
-  CHECK_CUDA_INPUT(dz);
-
-  int64_t count = z.numel();
-
-  AT_DISPATCH_FLOATING_TYPES(z.type(), "leaky_relu_backward_cuda", ([&] {
-    leaky_relu_backward_impl<scalar_t>(z.data<scalar_t>(), dz.data<scalar_t>(), slope, count);
-  }));
-}
-
-template<typename T>
-inline void elu_backward_impl(T *z, T *dz, int64_t count) {
-  // Create thrust pointers
-  thrust::device_ptr<T> th_z = thrust::device_pointer_cast(z);
-  thrust::device_ptr<T> th_dz = thrust::device_pointer_cast(dz);
-
-  auto stream = at::cuda::getCurrentCUDAStream();
-  thrust::transform_if(thrust::cuda::par.on(stream),
-                       th_dz, th_dz + count, th_z, th_z, th_dz,
-                       [] __device__ (const T& dz, const T& z) { return dz * (z + 1.); },
-                       [] __device__ (const T& z) { return z < 0; });
-  thrust::transform_if(thrust::cuda::par.on(stream),
-                       th_z, th_z + count, th_z,
-                       [] __device__ (const T& z) { return log1p(z); },
-                       [] __device__ (const T& z) { return z < 0; });
-}
-
-void elu_backward_cuda(at::Tensor z, at::Tensor dz) {
-  CHECK_CUDA_INPUT(z);
-  CHECK_CUDA_INPUT(dz);
-
-  int64_t count = z.numel();
-
-  AT_DISPATCH_FLOATING_TYPES(z.type(), "leaky_relu_backward_cuda", ([&] {
-    elu_backward_impl<scalar_t>(z.data<scalar_t>(), dz.data<scalar_t>(), count);
-  }));
-}

idm_masking/humanparsing/modules/src/inplace_abn_cuda_half.cu

@@ -1,275 +0,0 @@
-#include <ATen/ATen.h>
-
-#include <cuda_fp16.h>
-
-#include <vector>
-
-#include "utils/checks.h"
-#include "utils/cuda.cuh"
-#include "inplace_abn.h"
-
-#include <ATen/cuda/CUDAContext.h>
-
-// Operations for reduce
-struct SumOpH {
-  __device__ SumOpH(const half *t, int c, int s)
-      : tensor(t), chn(c), sp(s) {}
-  __device__ __forceinline__ float operator()(int batch, int plane, int n) {
-    return __half2float(tensor[(batch * chn + plane) * sp + n]);
-  }
-  const half *tensor;
-  const int chn;
-  const int sp;
-};
-
-struct VarOpH {
-  __device__ VarOpH(float m, const half *t, int c, int s)
-      : mean(m), tensor(t), chn(c), sp(s) {}
-  __device__ __forceinline__ float operator()(int batch, int plane, int n) {
-    const auto t = __half2float(tensor[(batch * chn + plane) * sp + n]);
-    return (t - mean) * (t - mean);
-  }
-  const float mean;
-  const half *tensor;
-  const int chn;
-  const int sp;
-};
-
-struct GradOpH {
-  __device__ GradOpH(float _weight, float _bias, const half *_z, const half *_dz, int c, int s)
-      : weight(_weight), bias(_bias), z(_z), dz(_dz), chn(c), sp(s) {}
-  __device__ __forceinline__ Pair<float> operator()(int batch, int plane, int n) {
-    float _y = (__half2float(z[(batch * chn + plane) * sp + n]) - bias) / weight;
-    float _dz = __half2float(dz[(batch * chn + plane) * sp + n]);
-    return Pair<float>(_dz, _y * _dz);
-  }
-  const float weight;
-  const float bias;
-  const half *z;
-  const half *dz;
-  const int chn;
-  const int sp;
-};
-
-/***********
- * mean_var
- ***********/
-
-__global__ void mean_var_kernel_h(const half *x, float *mean, float *var, int num, int chn, int sp) {
-  int plane = blockIdx.x;
-  float norm = 1.f / static_cast<float>(num * sp);
-
-  float _mean = reduce<float, SumOpH>(SumOpH(x, chn, sp), plane, num, sp) * norm;
-  __syncthreads();
-  float _var = reduce<float, VarOpH>(VarOpH(_mean, x, chn, sp), plane, num, sp) * norm;
-
-  if (threadIdx.x == 0) {
-    mean[plane] = _mean;
-    var[plane] = _var;
-  }
-}
-
-std::vector<at::Tensor> mean_var_cuda_h(at::Tensor x) {
-  CHECK_CUDA_INPUT(x);
-
-  // Extract dimensions
-  int64_t num, chn, sp;
-  get_dims(x, num, chn, sp);
-
-  // Prepare output tensors
-  auto mean = at::empty({chn},x.options().dtype(at::kFloat));
-  auto var = at::empty({chn},x.options().dtype(at::kFloat));
-
-  // Run kernel
-  dim3 blocks(chn);
-  dim3 threads(getNumThreads(sp));
-  auto stream = at::cuda::getCurrentCUDAStream();
-  mean_var_kernel_h<<<blocks, threads, 0, stream>>>(
-      reinterpret_cast<half*>(x.data<at::Half>()),
-      mean.data<float>(),
-      var.data<float>(),
-      num, chn, sp);
-
-  return {mean, var};
-}
-
-/**********
- * forward
- **********/
-
-__global__ void forward_kernel_h(half *x, const float *mean, const float *var, const float *weight, const float *bias,
-                                 bool affine, float eps, int num, int chn, int sp) {
-  int plane = blockIdx.x;
-
-  const float _mean = mean[plane];
-  const float _var = var[plane];
-  const float _weight = affine ? abs(weight[plane]) + eps : 1.f;
-  const float _bias = affine ? bias[plane] : 0.f;
-
-  const float mul = rsqrt(_var + eps) * _weight;
-
-  for (int batch = 0; batch < num; ++batch) {
-    for (int n = threadIdx.x; n < sp; n += blockDim.x) {
-      half *x_ptr = x + (batch * chn + plane) * sp + n;
-      float _x = __half2float(*x_ptr);
-      float _y = (_x - _mean) * mul + _bias;
-
-      *x_ptr = __float2half(_y);
-    }
-  }
-}
-
-at::Tensor forward_cuda_h(at::Tensor x, at::Tensor mean, at::Tensor var, at::Tensor weight, at::Tensor bias,
-                        bool affine, float eps) {
-  CHECK_CUDA_INPUT(x);
-  CHECK_CUDA_INPUT(mean);
-  CHECK_CUDA_INPUT(var);
-  CHECK_CUDA_INPUT(weight);
-  CHECK_CUDA_INPUT(bias);
-
-  // Extract dimensions
-  int64_t num, chn, sp;
-  get_dims(x, num, chn, sp);
-
-  // Run kernel
-  dim3 blocks(chn);
-  dim3 threads(getNumThreads(sp));
-  auto stream = at::cuda::getCurrentCUDAStream();
-  forward_kernel_h<<<blocks, threads, 0, stream>>>(
-      reinterpret_cast<half*>(x.data<at::Half>()),
-      mean.data<float>(),
-      var.data<float>(),
-      weight.data<float>(),
-      bias.data<float>(),
-      affine, eps, num, chn, sp);
-
-  return x;
-}
-
-__global__ void edz_eydz_kernel_h(const half *z, const half *dz, const float *weight, const float *bias,
-                                float *edz, float *eydz, bool affine, float eps, int num, int chn, int sp) {
-  int plane = blockIdx.x;
-
-  float _weight = affine ? abs(weight[plane]) + eps : 1.f;
-  float _bias = affine ? bias[plane] : 0.f;
-
-  Pair<float> res = reduce<Pair<float>, GradOpH>(GradOpH(_weight, _bias, z, dz, chn, sp), plane, num, sp);
-  __syncthreads();
-
-  if (threadIdx.x == 0) {
-    edz[plane] = res.v1;
-    eydz[plane] = res.v2;
-  }
-}
-
-std::vector<at::Tensor> edz_eydz_cuda_h(at::Tensor z, at::Tensor dz, at::Tensor weight, at::Tensor bias,
-                                      bool affine, float eps) {
-  CHECK_CUDA_INPUT(z);
-  CHECK_CUDA_INPUT(dz);
-  CHECK_CUDA_INPUT(weight);
-  CHECK_CUDA_INPUT(bias);
-
-  // Extract dimensions
-  int64_t num, chn, sp;
-  get_dims(z, num, chn, sp);
-
-  auto edz = at::empty({chn},z.options().dtype(at::kFloat));
-  auto eydz = at::empty({chn},z.options().dtype(at::kFloat));
-
-  // Run kernel
-  dim3 blocks(chn);
-  dim3 threads(getNumThreads(sp));
-  auto stream = at::cuda::getCurrentCUDAStream();
-  edz_eydz_kernel_h<<<blocks, threads, 0, stream>>>(
-        reinterpret_cast<half*>(z.data<at::Half>()),
-        reinterpret_cast<half*>(dz.data<at::Half>()),
-        weight.data<float>(),
-        bias.data<float>(),
-        edz.data<float>(),
-        eydz.data<float>(),
-        affine, eps, num, chn, sp);
- 
-  return {edz, eydz};
-}
-
-__global__ void backward_kernel_h(const half *z, const half *dz, const float *var, const float *weight, const float *bias, const float *edz,
-                                  const float *eydz, half *dx, bool affine, float eps, int num, int chn, int sp) {
-  int plane = blockIdx.x;
-
-  float _weight = affine ? abs(weight[plane]) + eps : 1.f;
-  float _bias = affine ? bias[plane] : 0.f;
-  float _var = var[plane];
-  float _edz = edz[plane];
-  float _eydz = eydz[plane];
-
-  float _mul = _weight * rsqrt(_var + eps);
-  float count = float(num * sp);
-
-  for (int batch = 0; batch < num; ++batch) {
-    for (int n = threadIdx.x; n < sp; n += blockDim.x) {
-      float _dz = __half2float(dz[(batch * chn + plane) * sp + n]);
-      float _y = (__half2float(z[(batch * chn + plane) * sp + n]) - _bias) / _weight;
-
-      dx[(batch * chn + plane) * sp + n] = __float2half((_dz - _edz / count - _y * _eydz / count) * _mul);
-    }
-  }
-}
-
-at::Tensor backward_cuda_h(at::Tensor z, at::Tensor dz, at::Tensor var, at::Tensor weight, at::Tensor bias,
-                                      at::Tensor edz, at::Tensor eydz, bool affine, float eps) {
-  CHECK_CUDA_INPUT(z);
-  CHECK_CUDA_INPUT(dz);
-  CHECK_CUDA_INPUT(var);
-  CHECK_CUDA_INPUT(weight);
-  CHECK_CUDA_INPUT(bias);
-  CHECK_CUDA_INPUT(edz);
-  CHECK_CUDA_INPUT(eydz);
-
-  // Extract dimensions
-  int64_t num, chn, sp;
-  get_dims(z, num, chn, sp);
-
-  auto dx = at::zeros_like(z);
-
-  // Run kernel
-  dim3 blocks(chn);
-  dim3 threads(getNumThreads(sp));
-  auto stream = at::cuda::getCurrentCUDAStream();
-  backward_kernel_h<<<blocks, threads, 0, stream>>>(
-        reinterpret_cast<half*>(z.data<at::Half>()),
-        reinterpret_cast<half*>(dz.data<at::Half>()),
-        var.data<float>(),
-        weight.data<float>(),
-        bias.data<float>(),
-        edz.data<float>(),
-        eydz.data<float>(),
-        reinterpret_cast<half*>(dx.data<at::Half>()),
-        affine, eps, num, chn, sp);
-
-  return dx;
-}
-
-__global__ void leaky_relu_backward_impl_h(half *z, half *dz, float slope, int64_t count) {
-  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < count;  i += blockDim.x * gridDim.x){
-    float _z = __half2float(z[i]);
-    if (_z < 0) {
-      dz[i] = __float2half(__half2float(dz[i]) * slope);
-      z[i] = __float2half(_z / slope);
-    }
-  }
-}
-
-void leaky_relu_backward_cuda_h(at::Tensor z, at::Tensor dz, float slope) {
-  CHECK_CUDA_INPUT(z);
-  CHECK_CUDA_INPUT(dz);
-
-  int64_t count = z.numel();
-  dim3 threads(getNumThreads(count));
-  dim3 blocks = (count + threads.x - 1) / threads.x;
-  auto stream = at::cuda::getCurrentCUDAStream();
-  leaky_relu_backward_impl_h<<<blocks, threads, 0, stream>>>(
-      reinterpret_cast<half*>(z.data<at::Half>()),
-      reinterpret_cast<half*>(dz.data<at::Half>()),
-      slope, count);
-}
-

idm_masking/humanparsing/modules/src/utils/checks.h

@@ -1,15 +0,0 @@
-#pragma once
-
-#include <ATen/ATen.h>
-
-// Define AT_CHECK for old version of ATen where the same function was called AT_ASSERT
-#ifndef AT_CHECK
-#define AT_CHECK AT_ASSERT
-#endif
-
-#define CHECK_CUDA(x) AT_CHECK((x).type().is_cuda(), #x " must be a CUDA tensor")
-#define CHECK_CPU(x) AT_CHECK(!(x).type().is_cuda(), #x " must be a CPU tensor")
-#define CHECK_CONTIGUOUS(x) AT_CHECK((x).is_contiguous(), #x " must be contiguous")
-
-#define CHECK_CUDA_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
-#define CHECK_CPU_INPUT(x) CHECK_CPU(x); CHECK_CONTIGUOUS(x)

idm_masking/humanparsing/modules/src/utils/common.h

@@ -1,49 +0,0 @@
-#pragma once
-
-#include <ATen/ATen.h>
-
-/*
- * Functions to share code between CPU and GPU
- */
-
-#ifdef __CUDACC__
-// CUDA versions
-
-#define HOST_DEVICE __host__ __device__
-#define INLINE_HOST_DEVICE __host__ __device__ inline
-#define FLOOR(x) floor(x)
-
-#if __CUDA_ARCH__ >= 600
-// Recent compute capabilities have block-level atomicAdd for all data types, so we use that
-#define ACCUM(x,y) atomicAdd_block(&(x),(y))
-#else
-// Older architectures don't have block-level atomicAdd, nor atomicAdd for doubles, so we defer to atomicAdd for float
-// and use the known atomicCAS-based implementation for double
-template<typename data_t>
-__device__ inline data_t atomic_add(data_t *address, data_t val) {
-  return atomicAdd(address, val);
-}
-
-template<>
-__device__ inline double atomic_add(double *address, double val) {
-  unsigned long long int* address_as_ull = (unsigned long long int*)address;
-  unsigned long long int old = *address_as_ull, assumed;
-  do {
-    assumed = old;
-    old = atomicCAS(address_as_ull, assumed, __double_as_longlong(val + __longlong_as_double(assumed)));
-  } while (assumed != old);
-  return __longlong_as_double(old);
-}
-
-#define ACCUM(x,y) atomic_add(&(x),(y))
-#endif // #if __CUDA_ARCH__ >= 600
-
-#else
-// CPU versions
-
-#define HOST_DEVICE
-#define INLINE_HOST_DEVICE inline
-#define FLOOR(x) std::floor(x)
-#define ACCUM(x,y) (x) += (y)
-
-#endif // #ifdef __CUDACC__

idm_masking/humanparsing/modules/src/utils/cuda.cuh

@@ -1,71 +0,0 @@
-#pragma once
-
-/*
- * General settings and functions
- */
-const int WARP_SIZE = 32;
-const int MAX_BLOCK_SIZE = 1024;
-
-static int getNumThreads(int nElem) {
-  int threadSizes[6] = {32, 64, 128, 256, 512, MAX_BLOCK_SIZE};
-  for (int i = 0; i < 6; ++i) {
-    if (nElem <= threadSizes[i]) {
-      return threadSizes[i];
-    }
-  }
-  return MAX_BLOCK_SIZE;
-}
-
-/*
- * Reduction utilities
- */
-template <typename T>
-__device__ __forceinline__ T WARP_SHFL_XOR(T value, int laneMask, int width = warpSize,
-                                           unsigned int mask = 0xffffffff) {
-#if CUDART_VERSION >= 9000
-  return __shfl_xor_sync(mask, value, laneMask, width);
-#else
-  return __shfl_xor(value, laneMask, width);
-#endif
-}
-
-__device__ __forceinline__ int getMSB(int val) { return 31 - __clz(val); }
-
-template<typename T>
-struct Pair {
-  T v1, v2;
-  __device__ Pair() {}
-  __device__ Pair(T _v1, T _v2) : v1(_v1), v2(_v2) {}
-  __device__ Pair(T v) : v1(v), v2(v) {}
-  __device__ Pair(int v) : v1(v), v2(v) {}
-  __device__ Pair &operator+=(const Pair<T> &a) {
-    v1 += a.v1;
-    v2 += a.v2;
-    return *this;
-  }
-};
-
-template<typename T>
-static __device__ __forceinline__ T warpSum(T val) {
-#if __CUDA_ARCH__ >= 300
-  for (int i = 0; i < getMSB(WARP_SIZE); ++i) {
-    val += WARP_SHFL_XOR(val, 1 << i, WARP_SIZE);
-  }
-#else
-  __shared__ T values[MAX_BLOCK_SIZE];
-  values[threadIdx.x] = val;
-  __threadfence_block();
-  const int base = (threadIdx.x / WARP_SIZE) * WARP_SIZE;
-  for (int i = 1; i < WARP_SIZE; i++) {
-    val += values[base + ((i + threadIdx.x) % WARP_SIZE)];
-  }
-#endif
-  return val;
-}
-
-template<typename T>
-static __device__ __forceinline__ Pair<T> warpSum(Pair<T> value) {
-  value.v1 = warpSum(value.v1);
-  value.v2 = warpSum(value.v2);
-  return value;
-}

idm_masking/humanparsing/networks/AugmentCE2P.py

@@ -1,388 +0,0 @@
-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-"""
-@Author  :   Peike Li
-@Contact :   peike.li@yahoo.com
-@File    :   AugmentCE2P.py
-@Time    :   8/4/19 3:35 PM
-@Desc    :
-@License :   This source code is licensed under the license found in the
-             LICENSE file in the root directory of this source tree.
-"""
-
-import functools
-import pdb
-
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-# Note here we adopt the InplaceABNSync implementation from https://github.com/mapillary/inplace_abn
-# By default, the InplaceABNSync module contains a BatchNorm Layer and a LeakyReLu layer
-from modules import InPlaceABNSync
-import numpy as np
-
-BatchNorm2d = functools.partial(InPlaceABNSync, activation='none')
-
-affine_par = True
-
-pretrained_settings = {
-    'resnet101': {
-        'imagenet': {
-            'input_space': 'BGR',
-            'input_size': [3, 224, 224],
-            'input_range': [0, 1],
-            'mean': [0.406, 0.456, 0.485],
-            'std': [0.225, 0.224, 0.229],
-            'num_classes': 1000
-        }
-    },
-}
-
-
-def conv3x3(in_planes, out_planes, stride=1):
-    "3x3 convolution with padding"
-    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
-                     padding=1, bias=False)
-
-
-class Bottleneck(nn.Module):
-    expansion = 4
-
-    def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, fist_dilation=1, multi_grid=1):
-        super(Bottleneck, self).__init__()
-        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
-        self.bn1 = BatchNorm2d(planes)
-        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
-                               padding=dilation * multi_grid, dilation=dilation * multi_grid, bias=False)
-        self.bn2 = BatchNorm2d(planes)
-        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
-        self.bn3 = BatchNorm2d(planes * 4)
-        self.relu = nn.ReLU(inplace=False)
-        self.relu_inplace = nn.ReLU(inplace=True)
-        self.downsample = downsample
-        self.dilation = dilation
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-        out = self.relu(out)
-
-        out = self.conv3(out)
-        out = self.bn3(out)
-
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out = out + residual
-        out = self.relu_inplace(out)
-
-        return out
-
-
-class CostomAdaptiveAvgPool2D(nn.Module):
-
-    def __init__(self, output_size):
-
-        super(CostomAdaptiveAvgPool2D, self).__init__()
-
-        self.output_size = output_size
-
-    def forward(self, x):
-
-        H_in, W_in = x.shape[-2:]
-        H_out, W_out = self.output_size
-
-        out_i = []
-        for i in range(H_out):
-            out_j = []
-            for j in range(W_out):
-                hs = int(np.floor(i * H_in / H_out))
-                he = int(np.ceil((i + 1) * H_in / H_out))
-
-                ws = int(np.floor(j * W_in / W_out))
-                we = int(np.ceil((j + 1) * W_in / W_out))
-
-                # print(hs, he, ws, we)
-                kernel_size = [he - hs, we - ws]
-
-                out = F.avg_pool2d(x[:, :, hs:he, ws:we], kernel_size)
-                out_j.append(out)
-
-            out_j = torch.concat(out_j, -1)
-            out_i.append(out_j)
-
-        out_i = torch.concat(out_i, -2)
-        return out_i
-
-
-class PSPModule(nn.Module):
-    """
-    Reference:
-        Zhao, Hengshuang, et al. *"Pyramid scene parsing network."*
-    """
-
-    def __init__(self, features, out_features=512, sizes=(1, 2, 3, 6)):
-        super(PSPModule, self).__init__()
-
-        self.stages = []
-        tmp = []
-        for size in sizes:
-            if size == 3 or size == 6:
-                tmp.append(self._make_stage_custom(features, out_features, size))
-            else:
-                tmp.append(self._make_stage(features, out_features, size))
-        self.stages = nn.ModuleList(tmp)
-        # self.stages = nn.ModuleList([self._make_stage(features, out_features, size) for size in sizes])
-        self.bottleneck = nn.Sequential(
-            nn.Conv2d(features + len(sizes) * out_features, out_features, kernel_size=3, padding=1, dilation=1,
-                      bias=False),
-            InPlaceABNSync(out_features),
-        )
-
-    def _make_stage(self, features, out_features, size):
-        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
-        conv = nn.Conv2d(features, out_features, kernel_size=1, bias=False)
-        bn = InPlaceABNSync(out_features)
-        return nn.Sequential(prior, conv, bn)
-
-    def _make_stage_custom(self, features, out_features, size):
-        prior = CostomAdaptiveAvgPool2D(output_size=(size, size))
-        conv = nn.Conv2d(features, out_features, kernel_size=1, bias=False)
-        bn = InPlaceABNSync(out_features)
-        return nn.Sequential(prior, conv, bn)
-
-    def forward(self, feats):
-        h, w = feats.size(2), feats.size(3)
-        priors = [F.interpolate(input=stage(feats), size=(h, w), mode='bilinear', align_corners=True) for stage in
-                  self.stages] + [feats]
-        bottle = self.bottleneck(torch.cat(priors, 1))
-        return bottle
-
-
-class ASPPModule(nn.Module):
-    """
-    Reference: 
-        Chen, Liang-Chieh, et al. *"Rethinking Atrous Convolution for Semantic Image Segmentation."*
-    """
-
-    def __init__(self, features, inner_features=256, out_features=512, dilations=(12, 24, 36)):
-        super(ASPPModule, self).__init__()
-
-        self.conv1 = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
-                                   nn.Conv2d(features, inner_features, kernel_size=1, padding=0, dilation=1,
-                                             bias=False),
-                                   InPlaceABNSync(inner_features))
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(features, inner_features, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(inner_features))
-        self.conv3 = nn.Sequential(
-            nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[0], dilation=dilations[0], bias=False),
-            InPlaceABNSync(inner_features))
-        self.conv4 = nn.Sequential(
-            nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[1], dilation=dilations[1], bias=False),
-            InPlaceABNSync(inner_features))
-        self.conv5 = nn.Sequential(
-            nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[2], dilation=dilations[2], bias=False),
-            InPlaceABNSync(inner_features))
-
-        self.bottleneck = nn.Sequential(
-            nn.Conv2d(inner_features * 5, out_features, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(out_features),
-            nn.Dropout2d(0.1)
-        )
-
-    def forward(self, x):
-        _, _, h, w = x.size()
-
-        feat1 = F.interpolate(self.conv1(x), size=(h, w), mode='bilinear', align_corners=True)
-
-        feat2 = self.conv2(x)
-        feat3 = self.conv3(x)
-        feat4 = self.conv4(x)
-        feat5 = self.conv5(x)
-        out = torch.cat((feat1, feat2, feat3, feat4, feat5), 1)
-
-        bottle = self.bottleneck(out)
-        return bottle
-
-
-class Edge_Module(nn.Module):
-    """
-    Edge Learning Branch
-    """
-
-    def __init__(self, in_fea=[256, 512, 1024], mid_fea=256, out_fea=2):
-        super(Edge_Module, self).__init__()
-
-        self.conv1 = nn.Sequential(
-            nn.Conv2d(in_fea[0], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(mid_fea)
-        )
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(in_fea[1], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(mid_fea)
-        )
-        self.conv3 = nn.Sequential(
-            nn.Conv2d(in_fea[2], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(mid_fea)
-        )
-        self.conv4 = nn.Conv2d(mid_fea, out_fea, kernel_size=3, padding=1, dilation=1, bias=True)
-        self.conv5 = nn.Conv2d(out_fea * 3, out_fea, kernel_size=1, padding=0, dilation=1, bias=True)
-
-    def forward(self, x1, x2, x3):
-        _, _, h, w = x1.size()
-
-        edge1_fea = self.conv1(x1)
-        edge1 = self.conv4(edge1_fea)
-        edge2_fea = self.conv2(x2)
-        edge2 = self.conv4(edge2_fea)
-        edge3_fea = self.conv3(x3)
-        edge3 = self.conv4(edge3_fea)
-
-        edge2_fea = F.interpolate(edge2_fea, size=(h, w), mode='bilinear', align_corners=True)
-        edge3_fea = F.interpolate(edge3_fea, size=(h, w), mode='bilinear', align_corners=True)
-        edge2 = F.interpolate(edge2, size=(h, w), mode='bilinear', align_corners=True)
-        edge3 = F.interpolate(edge3, size=(h, w), mode='bilinear', align_corners=True)
-
-        edge = torch.cat([edge1, edge2, edge3], dim=1)
-        edge_fea = torch.cat([edge1_fea, edge2_fea, edge3_fea], dim=1)
-        edge = self.conv5(edge)
-
-        return edge, edge_fea
-
-
-class Decoder_Module(nn.Module):
-    """
-    Parsing Branch Decoder Module.
-    """
-
-    def __init__(self, num_classes):
-        super(Decoder_Module, self).__init__()
-        self.conv1 = nn.Sequential(
-            nn.Conv2d(512, 256, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(256)
-        )
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(256, 48, kernel_size=1, stride=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(48)
-        )
-        self.conv3 = nn.Sequential(
-            nn.Conv2d(304, 256, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(256),
-            nn.Conv2d(256, 256, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(256)
-        )
-
-        self.conv4 = nn.Conv2d(256, num_classes, kernel_size=1, padding=0, dilation=1, bias=True)
-
-    def forward(self, xt, xl):
-        _, _, h, w = xl.size()
-        xt = F.interpolate(self.conv1(xt), size=(h, w), mode='bilinear', align_corners=True)
-        xl = self.conv2(xl)
-        x = torch.cat([xt, xl], dim=1)
-        x = self.conv3(x)
-        seg = self.conv4(x)
-        return seg, x
-
-
-class ResNet(nn.Module):
-    def __init__(self, block, layers, num_classes):
-        self.inplanes = 128
-        super(ResNet, self).__init__()
-        self.conv1 = conv3x3(3, 64, stride=2)
-        self.bn1 = BatchNorm2d(64)
-        self.relu1 = nn.ReLU(inplace=False)
-        self.conv2 = conv3x3(64, 64)
-        self.bn2 = BatchNorm2d(64)
-        self.relu2 = nn.ReLU(inplace=False)
-        self.conv3 = conv3x3(64, 128)
-        self.bn3 = BatchNorm2d(128)
-        self.relu3 = nn.ReLU(inplace=False)
-
-        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
-
-        self.layer1 = self._make_layer(block, 64, layers[0])
-        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
-        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
-        self.layer4 = self._make_layer(block, 512, layers[3], stride=1, dilation=2, multi_grid=(1, 1, 1))
-
-        self.context_encoding = PSPModule(2048, 512)
-
-        self.edge = Edge_Module()
-        self.decoder = Decoder_Module(num_classes)
-
-        self.fushion = nn.Sequential(
-            nn.Conv2d(1024, 256, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(256),
-            nn.Dropout2d(0.1),
-            nn.Conv2d(256, num_classes, kernel_size=1, padding=0, dilation=1, bias=True)
-        )
-
-    def _make_layer(self, block, planes, blocks, stride=1, dilation=1, multi_grid=1):
-        downsample = None
-        if stride != 1 or self.inplanes != planes * block.expansion:
-            downsample = nn.Sequential(
-                nn.Conv2d(self.inplanes, planes * block.expansion,
-                          kernel_size=1, stride=stride, bias=False),
-                BatchNorm2d(planes * block.expansion, affine=affine_par))
-
-        layers = []
-        generate_multi_grid = lambda index, grids: grids[index % len(grids)] if isinstance(grids, tuple) else 1
-        layers.append(block(self.inplanes, planes, stride, dilation=dilation, downsample=downsample,
-                            multi_grid=generate_multi_grid(0, multi_grid)))
-        self.inplanes = planes * block.expansion
-        for i in range(1, blocks):
-            layers.append(
-                block(self.inplanes, planes, dilation=dilation, multi_grid=generate_multi_grid(i, multi_grid)))
-
-        return nn.Sequential(*layers)
-
-    def forward(self, x):
-        x = self.relu1(self.bn1(self.conv1(x)))
-        x = self.relu2(self.bn2(self.conv2(x)))
-        x = self.relu3(self.bn3(self.conv3(x)))
-        x = self.maxpool(x)
-        x2 = self.layer1(x)
-        x3 = self.layer2(x2)
-        x4 = self.layer3(x3)
-        x5 = self.layer4(x4)
-        x = self.context_encoding(x5)
-        parsing_result, parsing_fea = self.decoder(x, x2)
-        # Edge Branch
-        edge_result, edge_fea = self.edge(x2, x3, x4)
-        # Fusion Branch
-        x = torch.cat([parsing_fea, edge_fea], dim=1)
-        fusion_result = self.fushion(x)
-        return [[parsing_result, fusion_result], edge_result]
-
-
-def initialize_pretrained_model(model, settings, pretrained='./models/resnet101-imagenet.pth'):
-    model.input_space = settings['input_space']
-    model.input_size = settings['input_size']
-    model.input_range = settings['input_range']
-    model.mean = settings['mean']
-    model.std = settings['std']
-
-    if pretrained is not None:
-        saved_state_dict = torch.load(pretrained)
-        new_params = model.state_dict().copy()
-        for i in saved_state_dict:
-            i_parts = i.split('.')
-            if not i_parts[0] == 'fc':
-                new_params['.'.join(i_parts[0:])] = saved_state_dict[i]
-        model.load_state_dict(new_params)
-
-
-def resnet101(num_classes=20, pretrained='./models/resnet101-imagenet.pth'):
-    model = ResNet(Bottleneck, [3, 4, 23, 3], num_classes)
-    settings = pretrained_settings['resnet101']['imagenet']
-    initialize_pretrained_model(model, settings, pretrained)
-    return model

idm_masking/humanparsing/networks/__init__.py

@@ -1,12 +0,0 @@
-from __future__ import absolute_import
-from networks.AugmentCE2P import resnet101
-
-__factory = {
-    'resnet101': resnet101,
-}
-
-
-def init_model(name, *args, **kwargs):
-    if name not in __factory.keys():
-        raise KeyError("Unknown model arch: {}".format(name))
-    return __factory[name](*args, **kwargs)

idm_masking/humanparsing/networks/backbone/mobilenetv2.py

@@ -1,156 +0,0 @@
-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-"""
-@Author  :   Peike Li
-@Contact :   peike.li@yahoo.com
-@File    :   mobilenetv2.py
-@Time    :   8/4/19 3:35 PM
-@Desc    :   
-@License :   This source code is licensed under the license found in the 
-             LICENSE file in the root directory of this source tree.
-"""
-
-import torch.nn as nn
-import math
-import functools
-
-from modules import InPlaceABN, InPlaceABNSync
-
-BatchNorm2d = functools.partial(InPlaceABNSync, activation='none')
-
-__all__ = ['mobilenetv2']
-
-
-def conv_bn(inp, oup, stride):
-    return nn.Sequential(
-        nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
-        BatchNorm2d(oup),
-        nn.ReLU6(inplace=True)
-    )
-
-
-def conv_1x1_bn(inp, oup):
-    return nn.Sequential(
-        nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
-        BatchNorm2d(oup),
-        nn.ReLU6(inplace=True)
-    )
-
-
-class InvertedResidual(nn.Module):
-    def __init__(self, inp, oup, stride, expand_ratio):
-        super(InvertedResidual, self).__init__()
-        self.stride = stride
-        assert stride in [1, 2]
-
-        hidden_dim = round(inp * expand_ratio)
-        self.use_res_connect = self.stride == 1 and inp == oup
-
-        if expand_ratio == 1:
-            self.conv = nn.Sequential(
-                # dw
-                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
-                BatchNorm2d(hidden_dim),
-                nn.ReLU6(inplace=True),
-                # pw-linear
-                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
-                BatchNorm2d(oup),
-            )
-        else:
-            self.conv = nn.Sequential(
-                # pw
-                nn.Conv2d(inp, hidden_dim, 1, 1, 0, bias=False),
-                BatchNorm2d(hidden_dim),
-                nn.ReLU6(inplace=True),
-                # dw
-                nn.Conv2d(hidden_dim, hidden_dim, 3, stride, 1, groups=hidden_dim, bias=False),
-                BatchNorm2d(hidden_dim),
-                nn.ReLU6(inplace=True),
-                # pw-linear
-                nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
-                BatchNorm2d(oup),
-            )
-
-    def forward(self, x):
-        if self.use_res_connect:
-            return x + self.conv(x)
-        else:
-            return self.conv(x)
-
-
-class MobileNetV2(nn.Module):
-    def __init__(self, n_class=1000, input_size=224, width_mult=1.):
-        super(MobileNetV2, self).__init__()
-        block = InvertedResidual
-        input_channel = 32
-        last_channel = 1280
-        interverted_residual_setting = [
-            # t, c, n, s
-            [1, 16, 1, 1],
-            [6, 24, 2, 2],  # layer 2
-            [6, 32, 3, 2],  # layer 3
-            [6, 64, 4, 2],
-            [6, 96, 3, 1],  # layer 4
-            [6, 160, 3, 2],
-            [6, 320, 1, 1],  # layer 5
-        ]
-
-        # building first layer
-        assert input_size % 32 == 0
-        input_channel = int(input_channel * width_mult)
-        self.last_channel = int(last_channel * width_mult) if width_mult > 1.0 else last_channel
-        self.features = [conv_bn(3, input_channel, 2)]
-        # building inverted residual blocks
-        for t, c, n, s in interverted_residual_setting:
-            output_channel = int(c * width_mult)
-            for i in range(n):
-                if i == 0:
-                    self.features.append(block(input_channel, output_channel, s, expand_ratio=t))
-                else:
-                    self.features.append(block(input_channel, output_channel, 1, expand_ratio=t))
-                input_channel = output_channel
-        # building last several layers
-        self.features.append(conv_1x1_bn(input_channel, self.last_channel))
-        # make it nn.Sequential
-        self.features = nn.Sequential(*self.features)
-
-        # building classifier
-        self.classifier = nn.Sequential(
-            nn.Dropout(0.2),
-            nn.Linear(self.last_channel, n_class),
-        )
-
-        self._initialize_weights()
-
-    def forward(self, x):
-        x = self.features(x)
-        x = x.mean(3).mean(2)
-        x = self.classifier(x)
-        return x
-
-    def _initialize_weights(self):
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
-                m.weight.data.normal_(0, math.sqrt(2. / n))
-                if m.bias is not None:
-                    m.bias.data.zero_()
-            elif isinstance(m, BatchNorm2d):
-                m.weight.data.fill_(1)
-                m.bias.data.zero_()
-            elif isinstance(m, nn.Linear):
-                n = m.weight.size(1)
-                m.weight.data.normal_(0, 0.01)
-                m.bias.data.zero_()
-
-
-def mobilenetv2(pretrained=False, **kwargs):
-    """Constructs a MobileNet_V2 model.
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
-    """
-    model = MobileNetV2(n_class=1000, **kwargs)
-    if pretrained:
-        model.load_state_dict(load_url(model_urls['mobilenetv2']), strict=False)
-    return model

idm_masking/humanparsing/networks/backbone/resnet.py

@@ -1,205 +0,0 @@
-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-"""
-@Author  :   Peike Li
-@Contact :   peike.li@yahoo.com
-@File    :   resnet.py
-@Time    :   8/4/19 3:35 PM
-@Desc    :   
-@License :   This source code is licensed under the license found in the 
-             LICENSE file in the root directory of this source tree.
-"""
-
-import functools
-import torch.nn as nn
-import math
-from torch.utils.model_zoo import load_url
-
-from modules import InPlaceABNSync
-
-BatchNorm2d = functools.partial(InPlaceABNSync, activation='none')
-
-__all__ = ['ResNet', 'resnet18', 'resnet50', 'resnet101']  # resnet101 is coming soon!
-
-model_urls = {
-    'resnet18': 'http://sceneparsing.csail.mit.edu/model/pretrained_resnet/resnet18-imagenet.pth',
-    'resnet50': 'http://sceneparsing.csail.mit.edu/model/pretrained_resnet/resnet50-imagenet.pth',
-    'resnet101': 'http://sceneparsing.csail.mit.edu/model/pretrained_resnet/resnet101-imagenet.pth'
-}
-
-
-def conv3x3(in_planes, out_planes, stride=1):
-    "3x3 convolution with padding"
-    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
-                     padding=1, bias=False)
-
-
-class BasicBlock(nn.Module):
-    expansion = 1
-
-    def __init__(self, inplanes, planes, stride=1, downsample=None):
-        super(BasicBlock, self).__init__()
-        self.conv1 = conv3x3(inplanes, planes, stride)
-        self.bn1 = BatchNorm2d(planes)
-        self.relu = nn.ReLU(inplace=True)
-        self.conv2 = conv3x3(planes, planes)
-        self.bn2 = BatchNorm2d(planes)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out += residual
-        out = self.relu(out)
-
-        return out
-
-
-class Bottleneck(nn.Module):
-    expansion = 4
-
-    def __init__(self, inplanes, planes, stride=1, downsample=None):
-        super(Bottleneck, self).__init__()
-        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
-        self.bn1 = BatchNorm2d(planes)
-        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
-                               padding=1, bias=False)
-        self.bn2 = BatchNorm2d(planes)
-        self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
-        self.bn3 = BatchNorm2d(planes * 4)
-        self.relu = nn.ReLU(inplace=True)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-        out = self.relu(out)
-
-        out = self.conv3(out)
-        out = self.bn3(out)
-
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out += residual
-        out = self.relu(out)
-
-        return out
-
-
-class ResNet(nn.Module):
-
-    def __init__(self, block, layers, num_classes=1000):
-        self.inplanes = 128
-        super(ResNet, self).__init__()
-        self.conv1 = conv3x3(3, 64, stride=2)
-        self.bn1 = BatchNorm2d(64)
-        self.relu1 = nn.ReLU(inplace=True)
-        self.conv2 = conv3x3(64, 64)
-        self.bn2 = BatchNorm2d(64)
-        self.relu2 = nn.ReLU(inplace=True)
-        self.conv3 = conv3x3(64, 128)
-        self.bn3 = BatchNorm2d(128)
-        self.relu3 = nn.ReLU(inplace=True)
-        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
-
-        self.layer1 = self._make_layer(block, 64, layers[0])
-        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
-        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
-        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
-        self.avgpool = nn.AvgPool2d(7, stride=1)
-        self.fc = nn.Linear(512 * block.expansion, num_classes)
-
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
-                m.weight.data.normal_(0, math.sqrt(2. / n))
-            elif isinstance(m, BatchNorm2d):
-                m.weight.data.fill_(1)
-                m.bias.data.zero_()
-
-    def _make_layer(self, block, planes, blocks, stride=1):
-        downsample = None
-        if stride != 1 or self.inplanes != planes * block.expansion:
-            downsample = nn.Sequential(
-                nn.Conv2d(self.inplanes, planes * block.expansion,
-                          kernel_size=1, stride=stride, bias=False),
-                BatchNorm2d(planes * block.expansion),
-            )
-
-        layers = []
-        layers.append(block(self.inplanes, planes, stride, downsample))
-        self.inplanes = planes * block.expansion
-        for i in range(1, blocks):
-            layers.append(block(self.inplanes, planes))
-
-        return nn.Sequential(*layers)
-
-    def forward(self, x):
-        x = self.relu1(self.bn1(self.conv1(x)))
-        x = self.relu2(self.bn2(self.conv2(x)))
-        x = self.relu3(self.bn3(self.conv3(x)))
-        x = self.maxpool(x)
-
-        x = self.layer1(x)
-        x = self.layer2(x)
-        x = self.layer3(x)
-        x = self.layer4(x)
-
-        x = self.avgpool(x)
-        x = x.view(x.size(0), -1)
-        x = self.fc(x)
-
-        return x
-
-
-def resnet18(pretrained=False, **kwargs):
-    """Constructs a ResNet-18 model.
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
-    """
-    model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
-    if pretrained:
-        model.load_state_dict(load_url(model_urls['resnet18']))
-    return model
-
-
-def resnet50(pretrained=False, **kwargs):
-    """Constructs a ResNet-50 model.
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
-    """
-    model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
-    if pretrained:
-        model.load_state_dict(load_url(model_urls['resnet50']), strict=False)
-    return model
-
-
-def resnet101(pretrained=False, **kwargs):
-    """Constructs a ResNet-101 model.
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on ImageNet
-    """
-    model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
-    if pretrained:
-        model.load_state_dict(load_url(model_urls['resnet101']), strict=False)
-    return model

idm_masking/humanparsing/networks/backbone/resnext.py

@@ -1,149 +0,0 @@
-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-"""
-@Author  :   Peike Li
-@Contact :   peike.li@yahoo.com
-@File    :   resnext.py.py
-@Time    :   8/11/19 8:58 PM
-@Desc    :   
-@License :   This source code is licensed under the license found in the 
-             LICENSE file in the root directory of this source tree.
-"""
-import functools
-import torch.nn as nn
-import math
-from torch.utils.model_zoo import load_url
-
-from modules import InPlaceABNSync
-
-BatchNorm2d = functools.partial(InPlaceABNSync, activation='none')
-
-__all__ = ['ResNeXt', 'resnext101']  # support resnext 101
-
-model_urls = {
-    'resnext50': 'http://sceneparsing.csail.mit.edu/model/pretrained_resnet/resnext50-imagenet.pth',
-    'resnext101': 'http://sceneparsing.csail.mit.edu/model/pretrained_resnet/resnext101-imagenet.pth'
-}
-
-
-def conv3x3(in_planes, out_planes, stride=1):
-    "3x3 convolution with padding"
-    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
-                     padding=1, bias=False)
-
-
-class GroupBottleneck(nn.Module):
-    expansion = 2
-
-    def __init__(self, inplanes, planes, stride=1, groups=1, downsample=None):
-        super(GroupBottleneck, self).__init__()
-        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
-        self.bn1 = BatchNorm2d(planes)
-        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
-                               padding=1, groups=groups, bias=False)
-        self.bn2 = BatchNorm2d(planes)
-        self.conv3 = nn.Conv2d(planes, planes * 2, kernel_size=1, bias=False)
-        self.bn3 = BatchNorm2d(planes * 2)
-        self.relu = nn.ReLU(inplace=True)
-        self.downsample = downsample
-        self.stride = stride
-
-    def forward(self, x):
-        residual = x
-
-        out = self.conv1(x)
-        out = self.bn1(out)
-        out = self.relu(out)
-
-        out = self.conv2(out)
-        out = self.bn2(out)
-        out = self.relu(out)
-
-        out = self.conv3(out)
-        out = self.bn3(out)
-
-        if self.downsample is not None:
-            residual = self.downsample(x)
-
-        out += residual
-        out = self.relu(out)
-
-        return out
-
-
-class ResNeXt(nn.Module):
-
-    def __init__(self, block, layers, groups=32, num_classes=1000):
-        self.inplanes = 128
-        super(ResNeXt, self).__init__()
-        self.conv1 = conv3x3(3, 64, stride=2)
-        self.bn1 = BatchNorm2d(64)
-        self.relu1 = nn.ReLU(inplace=True)
-        self.conv2 = conv3x3(64, 64)
-        self.bn2 = BatchNorm2d(64)
-        self.relu2 = nn.ReLU(inplace=True)
-        self.conv3 = conv3x3(64, 128)
-        self.bn3 = BatchNorm2d(128)
-        self.relu3 = nn.ReLU(inplace=True)
-        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
-
-        self.layer1 = self._make_layer(block, 128, layers[0], groups=groups)
-        self.layer2 = self._make_layer(block, 256, layers[1], stride=2, groups=groups)
-        self.layer3 = self._make_layer(block, 512, layers[2], stride=2, groups=groups)
-        self.layer4 = self._make_layer(block, 1024, layers[3], stride=2, groups=groups)
-        self.avgpool = nn.AvgPool2d(7, stride=1)
-        self.fc = nn.Linear(1024 * block.expansion, num_classes)
-
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels // m.groups
-                m.weight.data.normal_(0, math.sqrt(2. / n))
-            elif isinstance(m, BatchNorm2d):
-                m.weight.data.fill_(1)
-                m.bias.data.zero_()
-
-    def _make_layer(self, block, planes, blocks, stride=1, groups=1):
-        downsample = None
-        if stride != 1 or self.inplanes != planes * block.expansion:
-            downsample = nn.Sequential(
-                nn.Conv2d(self.inplanes, planes * block.expansion,
-                          kernel_size=1, stride=stride, bias=False),
-                BatchNorm2d(planes * block.expansion),
-            )
-
-        layers = []
-        layers.append(block(self.inplanes, planes, stride, groups, downsample))
-        self.inplanes = planes * block.expansion
-        for i in range(1, blocks):
-            layers.append(block(self.inplanes, planes, groups=groups))
-
-        return nn.Sequential(*layers)
-
-    def forward(self, x):
-        x = self.relu1(self.bn1(self.conv1(x)))
-        x = self.relu2(self.bn2(self.conv2(x)))
-        x = self.relu3(self.bn3(self.conv3(x)))
-        x = self.maxpool(x)
-
-        x = self.layer1(x)
-        x = self.layer2(x)
-        x = self.layer3(x)
-        x = self.layer4(x)
-
-        x = self.avgpool(x)
-        x = x.view(x.size(0), -1)
-        x = self.fc(x)
-
-        return x
-
-
-def resnext101(pretrained=False, **kwargs):
-    """Constructs a ResNet-101 model.
-    Args:
-        pretrained (bool): If True, returns a model pre-trained on Places
-    """
-    model = ResNeXt(GroupBottleneck, [3, 4, 23, 3], **kwargs)
-    if pretrained:
-        model.load_state_dict(load_url(model_urls['resnext101']), strict=False)
-    return model

idm_masking/humanparsing/networks/context_encoding/aspp.py

@@ -1,64 +0,0 @@
-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-"""
-@Author  :   Peike Li
-@Contact :   peike.li@yahoo.com
-@File    :   aspp.py
-@Time    :   8/4/19 3:36 PM
-@Desc    :   
-@License :   This source code is licensed under the license found in the 
-             LICENSE file in the root directory of this source tree.
-"""
-
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-
-from modules import InPlaceABNSync
-
-
-class ASPPModule(nn.Module):
-    """
-    Reference:
-        Chen, Liang-Chieh, et al. *"Rethinking Atrous Convolution for Semantic Image Segmentation."*
-    """
-    def __init__(self, features, out_features=512, inner_features=256, dilations=(12, 24, 36)):
-        super(ASPPModule, self).__init__()
-
-        self.conv1 = nn.Sequential(nn.AdaptiveAvgPool2d((1, 1)),
-                                   nn.Conv2d(features, inner_features, kernel_size=1, padding=0, dilation=1,
-                                             bias=False),
-                                   InPlaceABNSync(inner_features))
-        self.conv2 = nn.Sequential(
-            nn.Conv2d(features, inner_features, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(inner_features))
-        self.conv3 = nn.Sequential(
-            nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[0], dilation=dilations[0], bias=False),
-            InPlaceABNSync(inner_features))
-        self.conv4 = nn.Sequential(
-            nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[1], dilation=dilations[1], bias=False),
-            InPlaceABNSync(inner_features))
-        self.conv5 = nn.Sequential(
-            nn.Conv2d(features, inner_features, kernel_size=3, padding=dilations[2], dilation=dilations[2], bias=False),
-            InPlaceABNSync(inner_features))
-
-        self.bottleneck = nn.Sequential(
-            nn.Conv2d(inner_features * 5, out_features, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(out_features),
-            nn.Dropout2d(0.1)
-        )
-
-    def forward(self, x):
-        _, _, h, w = x.size()
-
-        feat1 = F.interpolate(self.conv1(x), size=(h, w), mode='bilinear', align_corners=True)
-
-        feat2 = self.conv2(x)
-        feat3 = self.conv3(x)
-        feat4 = self.conv4(x)
-        feat5 = self.conv5(x)
-        out = torch.cat((feat1, feat2, feat3, feat4, feat5), 1)
-
-        bottle = self.bottleneck(out)
-        return bottle

idm_masking/humanparsing/networks/context_encoding/ocnet.py

@@ -1,226 +0,0 @@
-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-"""
-@Author  :   Peike Li
-@Contact :   peike.li@yahoo.com
-@File    :   ocnet.py
-@Time    :   8/4/19 3:36 PM
-@Desc    :   
-@License :   This source code is licensed under the license found in the 
-             LICENSE file in the root directory of this source tree.
-"""
-
-import functools
-
-import torch
-import torch.nn as nn
-from torch.autograd import Variable
-from torch.nn import functional as F
-
-from modules import InPlaceABNSync
-BatchNorm2d = functools.partial(InPlaceABNSync, activation='none')
-
-
-class _SelfAttentionBlock(nn.Module):
-    '''
-    The basic implementation for self-attention block/non-local block
-    Input:
-        N X C X H X W
-    Parameters:
-        in_channels       : the dimension of the input feature map
-        key_channels      : the dimension after the key/query transform
-        value_channels    : the dimension after the value transform
-        scale             : choose the scale to downsample the input feature maps (save memory cost)
-    Return:
-        N X C X H X W
-        position-aware context features.(w/o concate or add with the input)
-    '''
-
-    def __init__(self, in_channels, key_channels, value_channels, out_channels=None, scale=1):
-        super(_SelfAttentionBlock, self).__init__()
-        self.scale = scale
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.key_channels = key_channels
-        self.value_channels = value_channels
-        if out_channels == None:
-            self.out_channels = in_channels
-        self.pool = nn.MaxPool2d(kernel_size=(scale, scale))
-        self.f_key = nn.Sequential(
-            nn.Conv2d(in_channels=self.in_channels, out_channels=self.key_channels,
-                      kernel_size=1, stride=1, padding=0),
-            InPlaceABNSync(self.key_channels),
-        )
-        self.f_query = self.f_key
-        self.f_value = nn.Conv2d(in_channels=self.in_channels, out_channels=self.value_channels,
-                                 kernel_size=1, stride=1, padding=0)
-        self.W = nn.Conv2d(in_channels=self.value_channels, out_channels=self.out_channels,
-                           kernel_size=1, stride=1, padding=0)
-        nn.init.constant(self.W.weight, 0)
-        nn.init.constant(self.W.bias, 0)
-
-    def forward(self, x):
-        batch_size, h, w = x.size(0), x.size(2), x.size(3)
-        if self.scale > 1:
-            x = self.pool(x)
-
-        value = self.f_value(x).view(batch_size, self.value_channels, -1)
-        value = value.permute(0, 2, 1)
-        query = self.f_query(x).view(batch_size, self.key_channels, -1)
-        query = query.permute(0, 2, 1)
-        key = self.f_key(x).view(batch_size, self.key_channels, -1)
-
-        sim_map = torch.matmul(query, key)
-        sim_map = (self.key_channels ** -.5) * sim_map
-        sim_map = F.softmax(sim_map, dim=-1)
-
-        context = torch.matmul(sim_map, value)
-        context = context.permute(0, 2, 1).contiguous()
-        context = context.view(batch_size, self.value_channels, *x.size()[2:])
-        context = self.W(context)
-        if self.scale > 1:
-            context = F.upsample(input=context, size=(h, w), mode='bilinear', align_corners=True)
-        return context
-
-
-class SelfAttentionBlock2D(_SelfAttentionBlock):
-    def __init__(self, in_channels, key_channels, value_channels, out_channels=None, scale=1):
-        super(SelfAttentionBlock2D, self).__init__(in_channels,
-                                                   key_channels,
-                                                   value_channels,
-                                                   out_channels,
-                                                   scale)
-
-
-class BaseOC_Module(nn.Module):
-    """
-    Implementation of the BaseOC module
-    Parameters:
-        in_features / out_features: the channels of the input / output feature maps.
-        dropout: we choose 0.05 as the default value.
-        size: you can apply multiple sizes. Here we only use one size.
-    Return:
-        features fused with Object context information.
-    """
-
-    def __init__(self, in_channels, out_channels, key_channels, value_channels, dropout, sizes=([1])):
-        super(BaseOC_Module, self).__init__()
-        self.stages = []
-        self.stages = nn.ModuleList(
-            [self._make_stage(in_channels, out_channels, key_channels, value_channels, size) for size in sizes])
-        self.conv_bn_dropout = nn.Sequential(
-            nn.Conv2d(2 * in_channels, out_channels, kernel_size=1, padding=0),
-            InPlaceABNSync(out_channels),
-            nn.Dropout2d(dropout)
-        )
-
-    def _make_stage(self, in_channels, output_channels, key_channels, value_channels, size):
-        return SelfAttentionBlock2D(in_channels,
-                                    key_channels,
-                                    value_channels,
-                                    output_channels,
-                                    size)
-
-    def forward(self, feats):
-        priors = [stage(feats) for stage in self.stages]
-        context = priors[0]
-        for i in range(1, len(priors)):
-            context += priors[i]
-        output = self.conv_bn_dropout(torch.cat([context, feats], 1))
-        return output
-
-
-class BaseOC_Context_Module(nn.Module):
-    """
-    Output only the context features.
-    Parameters:
-        in_features / out_features: the channels of the input / output feature maps.
-        dropout: specify the dropout ratio
-        fusion: We provide two different fusion method, "concat" or "add"
-        size: we find that directly learn the attention weights on even 1/8 feature maps is hard.
-    Return:
-        features after "concat" or "add"
-    """
-
-    def __init__(self, in_channels, out_channels, key_channels, value_channels, dropout, sizes=([1])):
-        super(BaseOC_Context_Module, self).__init__()
-        self.stages = []
-        self.stages = nn.ModuleList(
-            [self._make_stage(in_channels, out_channels, key_channels, value_channels, size) for size in sizes])
-        self.conv_bn_dropout = nn.Sequential(
-            nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0),
-            InPlaceABNSync(out_channels),
-        )
-
-    def _make_stage(self, in_channels, output_channels, key_channels, value_channels, size):
-        return SelfAttentionBlock2D(in_channels,
-                                    key_channels,
-                                    value_channels,
-                                    output_channels,
-                                    size)
-
-    def forward(self, feats):
-        priors = [stage(feats) for stage in self.stages]
-        context = priors[0]
-        for i in range(1, len(priors)):
-            context += priors[i]
-        output = self.conv_bn_dropout(context)
-        return output
-
-
-class ASP_OC_Module(nn.Module):
-    def __init__(self, features, out_features=256, dilations=(12, 24, 36)):
-        super(ASP_OC_Module, self).__init__()
-        self.context = nn.Sequential(nn.Conv2d(features, out_features, kernel_size=3, padding=1, dilation=1, bias=True),
-                                     InPlaceABNSync(out_features),
-                                     BaseOC_Context_Module(in_channels=out_features, out_channels=out_features,
-                                                           key_channels=out_features // 2, value_channels=out_features,
-                                                           dropout=0, sizes=([2])))
-        self.conv2 = nn.Sequential(nn.Conv2d(features, out_features, kernel_size=1, padding=0, dilation=1, bias=False),
-                                   InPlaceABNSync(out_features))
-        self.conv3 = nn.Sequential(
-            nn.Conv2d(features, out_features, kernel_size=3, padding=dilations[0], dilation=dilations[0], bias=False),
-            InPlaceABNSync(out_features))
-        self.conv4 = nn.Sequential(
-            nn.Conv2d(features, out_features, kernel_size=3, padding=dilations[1], dilation=dilations[1], bias=False),
-            InPlaceABNSync(out_features))
-        self.conv5 = nn.Sequential(
-            nn.Conv2d(features, out_features, kernel_size=3, padding=dilations[2], dilation=dilations[2], bias=False),
-            InPlaceABNSync(out_features))
-
-        self.conv_bn_dropout = nn.Sequential(
-            nn.Conv2d(out_features * 5, out_features, kernel_size=1, padding=0, dilation=1, bias=False),
-            InPlaceABNSync(out_features),
-            nn.Dropout2d(0.1)
-        )
-
-    def _cat_each(self, feat1, feat2, feat3, feat4, feat5):
-        assert (len(feat1) == len(feat2))
-        z = []
-        for i in range(len(feat1)):
-            z.append(torch.cat((feat1[i], feat2[i], feat3[i], feat4[i], feat5[i]), 1))
-        return z
-
-    def forward(self, x):
-        if isinstance(x, Variable):
-            _, _, h, w = x.size()
-        elif isinstance(x, tuple) or isinstance(x, list):
-            _, _, h, w = x[0].size()
-        else:
-            raise RuntimeError('unknown input type')
-
-        feat1 = self.context(x)
-        feat2 = self.conv2(x)
-        feat3 = self.conv3(x)
-        feat4 = self.conv4(x)
-        feat5 = self.conv5(x)
-
-        if isinstance(x, Variable):
-            out = torch.cat((feat1, feat2, feat3, feat4, feat5), 1)
-        elif isinstance(x, tuple) or isinstance(x, list):
-            out = self._cat_each(feat1, feat2, feat3, feat4, feat5)
-        else:
-            raise RuntimeError('unknown input type')
-        output = self.conv_bn_dropout(out)
-        return output

idm_masking/humanparsing/networks/context_encoding/psp.py

@@ -1,48 +0,0 @@
-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-"""
-@Author  :   Peike Li
-@Contact :   peike.li@yahoo.com
-@File    :   psp.py
-@Time    :   8/4/19 3:36 PM
-@Desc    :   
-@License :   This source code is licensed under the license found in the 
-             LICENSE file in the root directory of this source tree.
-"""
-
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-
-from modules import InPlaceABNSync
-
-
-class PSPModule(nn.Module):
-    """
-    Reference:
-        Zhao, Hengshuang, et al. *"Pyramid scene parsing network."*
-    """
-    def __init__(self, features, out_features=512, sizes=(1, 2, 3, 6)):
-        super(PSPModule, self).__init__()
-
-        self.stages = []
-        self.stages = nn.ModuleList([self._make_stage(features, out_features, size) for size in sizes])
-        self.bottleneck = nn.Sequential(
-            nn.Conv2d(features + len(sizes) * out_features, out_features, kernel_size=3, padding=1, dilation=1,
-                      bias=False),
-            InPlaceABNSync(out_features),
-        )
-
-    def _make_stage(self, features, out_features, size):
-        prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
-        conv = nn.Conv2d(features, out_features, kernel_size=1, bias=False)
-        bn = InPlaceABNSync(out_features)
-        return nn.Sequential(prior, conv, bn)
-
-    def forward(self, feats):
-        h, w = feats.size(2), feats.size(3)
-        priors = [F.interpolate(input=stage(feats), size=(h, w), mode='bilinear', align_corners=True) for stage in
-                  self.stages] + [feats]
-        bottle = self.bottleneck(torch.cat(priors, 1))
-        return bottle

idm_masking/humanparsing/parsing_api.py

@@ -1,188 +0,0 @@
-import pdb
-from pathlib import Path
-import sys
-
-PROJECT_ROOT = Path(__file__).absolute().parents[0].absolute()
-sys.path.insert(0, str(PROJECT_ROOT))
-import os
-import torch
-import numpy as np
-import cv2
-import torchvision.transforms as transforms
-from torch.utils.data import DataLoader
-from tools.simple_extractor_dataset import SimpleFolderDataset
-from tools.transforms import transform_logits
-from tqdm import tqdm
-from PIL import Image
-
-
-def get_palette(num_cls):
-    """ Returns the color map for visualizing the segmentation mask.
-    Args:
-        num_cls: Number of classes
-    Returns:
-        The color map
-    """
-    n = num_cls
-    palette = [0] * (n * 3)
-    for j in range(0, n):
-        lab = j
-        palette[j * 3 + 0] = 0
-        palette[j * 3 + 1] = 0
-        palette[j * 3 + 2] = 0
-        i = 0
-        while lab:
-            palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i))
-            palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i))
-            palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i))
-            i += 1
-            lab >>= 3
-    return palette
-
-
-def delete_irregular(logits_result):
-    parsing_result = np.argmax(logits_result, axis=2)
-    upper_cloth = np.where(parsing_result == 4, 255, 0)
-    contours, hierarchy = cv2.findContours(upper_cloth.astype(np.uint8),
-                                           cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_L1)
-    area = []
-    for i in range(len(contours)):
-        a = cv2.contourArea(contours[i], True)
-        area.append(abs(a))
-    if len(area) != 0:
-        top = area.index(max(area))
-        M = cv2.moments(contours[top])
-        cY = int(M["m01"] / M["m00"])
-
-    dresses = np.where(parsing_result == 7, 255, 0)
-    contours_dress, hierarchy_dress = cv2.findContours(dresses.astype(np.uint8),
-                                                       cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_L1)
-    area_dress = []
-    for j in range(len(contours_dress)):
-        a_d = cv2.contourArea(contours_dress[j], True)
-        area_dress.append(abs(a_d))
-    if len(area_dress) != 0:
-        top_dress = area_dress.index(max(area_dress))
-        M_dress = cv2.moments(contours_dress[top_dress])
-        cY_dress = int(M_dress["m01"] / M_dress["m00"])
-    wear_type = "dresses"
-    if len(area) != 0:
-        if len(area_dress) != 0 and cY_dress > cY:
-            irregular_list = np.array([4, 5, 6])
-            logits_result[:, :, irregular_list] = -1
-        else:
-            irregular_list = np.array([5, 6, 7, 8, 9, 10, 12, 13])
-            logits_result[:cY, :, irregular_list] = -1
-            wear_type = "cloth_pant"
-        parsing_result = np.argmax(logits_result, axis=2)
-    # pad border
-    parsing_result = np.pad(parsing_result, pad_width=1, mode='constant', constant_values=0)
-    return parsing_result, wear_type
-
-
-
-def hole_fill(img):
-    img_copy = img.copy()
-    mask = np.zeros((img.shape[0] + 2, img.shape[1] + 2), dtype=np.uint8)
-    cv2.floodFill(img, mask, (0, 0), 255)
-    img_inverse = cv2.bitwise_not(img)
-    dst = cv2.bitwise_or(img_copy, img_inverse)
-    return dst
-
-def refine_mask(mask):
-    contours, hierarchy = cv2.findContours(mask.astype(np.uint8),
-                                           cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_L1)
-    area = []
-    for j in range(len(contours)):
-        a_d = cv2.contourArea(contours[j], True)
-        area.append(abs(a_d))
-    refine_mask = np.zeros_like(mask).astype(np.uint8)
-    if len(area) != 0:
-        i = area.index(max(area))
-        cv2.drawContours(refine_mask, contours, i, color=255, thickness=-1)
-        # keep large area in skin case
-        for j in range(len(area)):
-          if j != i and area[i] > 2000:
-             cv2.drawContours(refine_mask, contours, j, color=255, thickness=-1)
-    return refine_mask
-
-def refine_hole(parsing_result_filled, parsing_result, arm_mask):
-    filled_hole = cv2.bitwise_and(np.where(parsing_result_filled == 4, 255, 0),
-                                  np.where(parsing_result != 4, 255, 0)) - arm_mask * 255
-    contours, hierarchy = cv2.findContours(filled_hole, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_L1)
-    refine_hole_mask = np.zeros_like(parsing_result).astype(np.uint8)
-    for i in range(len(contours)):
-        a = cv2.contourArea(contours[i], True)
-        # keep hole > 2000 pixels
-        if abs(a) > 2000:
-            cv2.drawContours(refine_hole_mask, contours, i, color=255, thickness=-1)
-    return refine_hole_mask + arm_mask
-
-def onnx_inference(session, lip_session, input_dir):
-    transform = transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.406, 0.456, 0.485], std=[0.225, 0.224, 0.229])
-    ])
-    dataset = SimpleFolderDataset(root=input_dir, input_size=[512, 512], transform=transform)
-    dataloader = DataLoader(dataset)
-    with torch.no_grad():
-        for _, batch in enumerate(tqdm(dataloader)):
-            image, meta = batch
-            c = meta['center'].numpy()[0]
-            s = meta['scale'].numpy()[0]
-            w = meta['width'].numpy()[0]
-            h = meta['height'].numpy()[0]
-            output = session.run(None, {"input.1": image.numpy().astype(np.float32)})
-            upsample = torch.nn.Upsample(size=[512, 512], mode='bilinear', align_corners=True)
-            upsample_output = upsample(torch.from_numpy(output[1][0]).unsqueeze(0))
-            upsample_output = upsample_output.squeeze()
-            upsample_output = upsample_output.permute(1, 2, 0)  # CHW -> HWC
-            logits_result = transform_logits(upsample_output.data.cpu().numpy(), c, s, w, h, input_size=[512, 512])
-            parsing_result = np.argmax(logits_result, axis=2)
-            parsing_result = np.pad(parsing_result, pad_width=1, mode='constant', constant_values=0)
-            # try holefilling the clothes part
-            arm_mask = (parsing_result == 14).astype(np.float32) \
-                       + (parsing_result == 15).astype(np.float32)
-            upper_cloth_mask = (parsing_result == 4).astype(np.float32) + arm_mask
-            img = np.where(upper_cloth_mask, 255, 0)
-            dst = hole_fill(img.astype(np.uint8))
-            parsing_result_filled = dst / 255 * 4
-            parsing_result_woarm = np.where(parsing_result_filled == 4, parsing_result_filled, parsing_result)
-            # add back arm and refined hole between arm and cloth
-            refine_hole_mask = refine_hole(parsing_result_filled.astype(np.uint8), parsing_result.astype(np.uint8),
-                                           arm_mask.astype(np.uint8))
-            parsing_result = np.where(refine_hole_mask, parsing_result, parsing_result_woarm)
-            # remove padding
-            parsing_result = parsing_result[1:-1, 1:-1]
-
-        dataset_lip = SimpleFolderDataset(root=input_dir, input_size=[473, 473], transform=transform)
-        dataloader_lip = DataLoader(dataset_lip)
-        with torch.no_grad():
-            for _, batch in enumerate(tqdm(dataloader_lip)):
-                image, meta = batch
-                c = meta['center'].numpy()[0]
-                s = meta['scale'].numpy()[0]
-                w = meta['width'].numpy()[0]
-                h = meta['height'].numpy()[0]
-
-                output_lip = lip_session.run(None, {"input.1": image.numpy().astype(np.float32)})
-                upsample = torch.nn.Upsample(size=[473, 473], mode='bilinear', align_corners=True)
-                upsample_output_lip = upsample(torch.from_numpy(output_lip[1][0]).unsqueeze(0))
-                upsample_output_lip = upsample_output_lip.squeeze()
-                upsample_output_lip = upsample_output_lip.permute(1, 2, 0)  # CHW -> HWC
-                logits_result_lip = transform_logits(upsample_output_lip.data.cpu().numpy(), c, s, w, h,
-                                                     input_size=[473, 473])
-                parsing_result_lip = np.argmax(logits_result_lip, axis=2)
-    # add neck parsing result
-    neck_mask = np.logical_and(np.logical_not((parsing_result_lip == 13).astype(np.float32)),
-                               (parsing_result == 11).astype(np.float32))
-    parsing_result = np.where(neck_mask, 18, parsing_result)
-    palette = get_palette(19)
-    output_img = Image.fromarray(np.asarray(parsing_result, dtype=np.uint8))
-    output_img.putpalette(palette)
-    face_mask = torch.from_numpy((parsing_result == 11).astype(np.float32))
-
-    return output_img, face_mask
-
-
-

idm_masking/humanparsing/run_parsing.py

@@ -1,29 +0,0 @@
-import pdb
-from pathlib import Path
-import sys
-import os
-import onnxruntime as ort
-PROJECT_ROOT = Path(__file__).absolute().parents[0].absolute()
-sys.path.insert(0, str(PROJECT_ROOT))
-from parsing_api import onnx_inference
-import torch
-
-
-class Parsing:
-    def __init__(self, gpu_id: int):
-        self.gpu_id = gpu_id
-        torch.cuda.set_device(gpu_id)
-        session_options = ort.SessionOptions()
-        session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-        session_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
-        session_options.add_session_config_entry('gpu_id', str(gpu_id))
-        self.session = ort.InferenceSession(os.path.join(Path(__file__).absolute().parents[2].absolute(), 'ckpt/humanparsing/parsing_atr.onnx'),
-                                            sess_options=session_options, providers=['CPUExecutionProvider'])
-        self.lip_session = ort.InferenceSession(os.path.join(Path(__file__).absolute().parents[2].absolute(), 'ckpt/humanparsing/parsing_lip.onnx'),
-                                                sess_options=session_options, providers=['CPUExecutionProvider'])
-        
-
-    def __call__(self, input_image):
-        # torch.cuda.set_device(self.gpu_id)
-        parsed_image, face_mask = onnx_inference(self.session, self.lip_session, input_image)
-        return parsed_image, face_mask

idm_masking/humanparsing/tools/datasets.py

@@ -1,201 +0,0 @@
-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-"""
-@Author  :   Peike Li
-@Contact :   peike.li@yahoo.com
-@File    :   datasets.py
-@Time    :   8/4/19 3:35 PM
-@Desc    :
-@License :   This source code is licensed under the license found in the
-             LICENSE file in the root directory of this source tree.
-"""
-
-import os
-import numpy as np
-import random
-import torch
-import cv2
-from torch.utils import data
-from .transforms import get_affine_transform
-
-
-class LIPDataSet(data.Dataset):
-    def __init__(self, root, dataset, crop_size=[473, 473], scale_factor=0.25,
-                 rotation_factor=30, ignore_label=255, transform=None):
-        self.root = root
-        self.aspect_ratio = crop_size[1] * 1.0 / crop_size[0]
-        self.crop_size = np.asarray(crop_size)
-        self.ignore_label = ignore_label
-        self.scale_factor = scale_factor
-        self.rotation_factor = rotation_factor
-        self.flip_prob = 0.5
-        self.transform = transform
-        self.dataset = dataset
-
-        list_path = os.path.join(self.root, self.dataset + '_id.txt')
-        train_list = [i_id.strip() for i_id in open(list_path)]
-
-        self.train_list = train_list
-        self.number_samples = len(self.train_list)
-
-    def __len__(self):
-        return self.number_samples
-
-    def _box2cs(self, box):
-        x, y, w, h = box[:4]
-        return self._xywh2cs(x, y, w, h)
-
-    def _xywh2cs(self, x, y, w, h):
-        center = np.zeros((2), dtype=np.float32)
-        center[0] = x + w * 0.5
-        center[1] = y + h * 0.5
-        if w > self.aspect_ratio * h:
-            h = w * 1.0 / self.aspect_ratio
-        elif w < self.aspect_ratio * h:
-            w = h * self.aspect_ratio
-        scale = np.array([w * 1.0, h * 1.0], dtype=np.float32)
-        return center, scale
-
-    def __getitem__(self, index):
-        train_item = self.train_list[index]
-
-        im_path = os.path.join(self.root, self.dataset + '_images', train_item + '.jpg')
-        parsing_anno_path = os.path.join(self.root, self.dataset + '_segmentations', train_item + '.png')
-
-        im = cv2.imread(im_path, cv2.IMREAD_COLOR)
-        h, w, _ = im.shape
-        parsing_anno = np.zeros((h, w), dtype=np.long)
-
-        # Get person center and scale
-        person_center, s = self._box2cs([0, 0, w - 1, h - 1])
-        r = 0
-
-        if self.dataset != 'test':
-            # Get pose annotation
-            parsing_anno = cv2.imread(parsing_anno_path, cv2.IMREAD_GRAYSCALE)
-            if self.dataset == 'train' or self.dataset == 'trainval':
-                sf = self.scale_factor
-                rf = self.rotation_factor
-                s = s * np.clip(np.random.randn() * sf + 1, 1 - sf, 1 + sf)
-                r = np.clip(np.random.randn() * rf, -rf * 2, rf * 2) if random.random() <= 0.6 else 0
-
-                if random.random() <= self.flip_prob:
-                    im = im[:, ::-1, :]
-                    parsing_anno = parsing_anno[:, ::-1]
-                    person_center[0] = im.shape[1] - person_center[0] - 1
-                    right_idx = [15, 17, 19]
-                    left_idx = [14, 16, 18]
-                    for i in range(0, 3):
-                        right_pos = np.where(parsing_anno == right_idx[i])
-                        left_pos = np.where(parsing_anno == left_idx[i])
-                        parsing_anno[right_pos[0], right_pos[1]] = left_idx[i]
-                        parsing_anno[left_pos[0], left_pos[1]] = right_idx[i]
-
-        trans = get_affine_transform(person_center, s, r, self.crop_size)
-        input = cv2.warpAffine(
-            im,
-            trans,
-            (int(self.crop_size[1]), int(self.crop_size[0])),
-            flags=cv2.INTER_LINEAR,
-            borderMode=cv2.BORDER_CONSTANT,
-            borderValue=(0, 0, 0))
-
-        if self.transform:
-            input = self.transform(input)
-
-        meta = {
-            'name': train_item,
-            'center': person_center,
-            'height': h,
-            'width': w,
-            'scale': s,
-            'rotation': r
-        }
-
-        if self.dataset == 'val' or self.dataset == 'test':
-            return input, meta
-        else:
-            label_parsing = cv2.warpAffine(
-                parsing_anno,
-                trans,
-                (int(self.crop_size[1]), int(self.crop_size[0])),
-                flags=cv2.INTER_NEAREST,
-                borderMode=cv2.BORDER_CONSTANT,
-                borderValue=(255))
-
-            label_parsing = torch.from_numpy(label_parsing)
-
-            return input, label_parsing, meta
-
-
-class LIPDataValSet(data.Dataset):
-    def __init__(self, root, dataset='val', crop_size=[473, 473], transform=None, flip=False):
-        self.root = root
-        self.crop_size = crop_size
-        self.transform = transform
-        self.flip = flip
-        self.dataset = dataset
-        self.root = root
-        self.aspect_ratio = crop_size[1] * 1.0 / crop_size[0]
-        self.crop_size = np.asarray(crop_size)
-
-        list_path = os.path.join(self.root, self.dataset + '_id.txt')
-        val_list = [i_id.strip() for i_id in open(list_path)]
-
-        self.val_list = val_list
-        self.number_samples = len(self.val_list)
-
-    def __len__(self):
-        return len(self.val_list)
-
-    def _box2cs(self, box):
-        x, y, w, h = box[:4]
-        return self._xywh2cs(x, y, w, h)
-
-    def _xywh2cs(self, x, y, w, h):
-        center = np.zeros((2), dtype=np.float32)
-        center[0] = x + w * 0.5
-        center[1] = y + h * 0.5
-        if w > self.aspect_ratio * h:
-            h = w * 1.0 / self.aspect_ratio
-        elif w < self.aspect_ratio * h:
-            w = h * self.aspect_ratio
-        scale = np.array([w * 1.0, h * 1.0], dtype=np.float32)
-
-        return center, scale
-
-    def __getitem__(self, index):
-        val_item = self.val_list[index]
-        # Load training image
-        im_path = os.path.join(self.root, self.dataset + '_images', val_item + '.jpg')
-        im = cv2.imread(im_path, cv2.IMREAD_COLOR)
-        h, w, _ = im.shape
-        # Get person center and scale
-        person_center, s = self._box2cs([0, 0, w - 1, h - 1])
-        r = 0
-        trans = get_affine_transform(person_center, s, r, self.crop_size)
-        input = cv2.warpAffine(
-            im,
-            trans,
-            (int(self.crop_size[1]), int(self.crop_size[0])),
-            flags=cv2.INTER_LINEAR,
-            borderMode=cv2.BORDER_CONSTANT,
-            borderValue=(0, 0, 0))
-        input = self.transform(input)
-        flip_input = input.flip(dims=[-1])
-        if self.flip:
-            batch_input_im = torch.stack([input, flip_input])
-        else:
-            batch_input_im = input
-
-        meta = {
-            'name': val_item,
-            'center': person_center,
-            'height': h,
-            'width': w,
-            'scale': s,
-            'rotation': r
-        }
-
-        return batch_input_im, meta

idm_masking/humanparsing/tools/encoding.py

@@ -1,188 +0,0 @@
-##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
-## Created by: Hang Zhang
-## ECE Department, Rutgers University
-## Email: zhang.hang@rutgers.edu
-## Copyright (c) 2017
-##
-## This source code is licensed under the MIT-style license found in the
-## LICENSE file in the root directory of this source tree
-##+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
-
-"""Encoding Data Parallel"""
-import threading
-import functools
-import torch
-from torch.autograd import Variable, Function
-import torch.cuda.comm as comm
-from torch.nn.parallel.data_parallel import DataParallel
-from torch.nn.parallel.parallel_apply import get_a_var
-from torch.nn.parallel._functions import ReduceAddCoalesced, Broadcast
-
-torch_ver = torch.__version__[:3]
-
-__all__ = ['allreduce', 'DataParallelModel', 'DataParallelCriterion', 'patch_replication_callback']
-
-def allreduce(*inputs):
-    """Cross GPU all reduce autograd operation for calculate mean and
-    variance in SyncBN.
-    """
-    return AllReduce.apply(*inputs)
-
-class AllReduce(Function):
-    @staticmethod
-    def forward(ctx, num_inputs, *inputs):
-        ctx.num_inputs = num_inputs
-        ctx.target_gpus = [inputs[i].get_device() for i in range(0, len(inputs), num_inputs)]
-        inputs = [inputs[i:i + num_inputs]
-                 for i in range(0, len(inputs), num_inputs)]
-        # sort before reduce sum
-        inputs = sorted(inputs, key=lambda i: i[0].get_device())
-        results = comm.reduce_add_coalesced(inputs, ctx.target_gpus[0])
-        outputs = comm.broadcast_coalesced(results, ctx.target_gpus)
-        return tuple([t for tensors in outputs for t in tensors])
-
-    @staticmethod
-    def backward(ctx, *inputs):
-        inputs = [i.data for i in inputs]
-        inputs = [inputs[i:i + ctx.num_inputs]
-                 for i in range(0, len(inputs), ctx.num_inputs)]
-        results = comm.reduce_add_coalesced(inputs, ctx.target_gpus[0])
-        outputs = comm.broadcast_coalesced(results, ctx.target_gpus)
-        return (None,) + tuple([Variable(t) for tensors in outputs for t in tensors])
-
-class Reduce(Function):
-    @staticmethod
-    def forward(ctx, *inputs):
-        ctx.target_gpus = [inputs[i].get_device() for i in range(len(inputs))]
-        inputs = sorted(inputs, key=lambda i: i.get_device())
-        return comm.reduce_add(inputs)
-
-    @staticmethod
-    def backward(ctx, gradOutput):
-        return Broadcast.apply(ctx.target_gpus, gradOutput)
-
-
-class DataParallelModel(DataParallel):
-    """Implements data parallelism at the module level.
-
-    This container parallelizes the application of the given module by
-    splitting the input across the specified devices by chunking in the
-    batch dimension.
-    In the forward pass, the module is replicated on each device,
-    and each replica handles a portion of the input. During the backwards pass, gradients from each replica are summed into the original module.
-    Note that the outputs are not gathered, please use compatible
-    :class:`encoding.parallel.DataParallelCriterion`.
-
-    The batch size should be larger than the number of GPUs used. It should
-    also be an integer multiple of the number of GPUs so that each chunk is
-    the same size (so that each GPU processes the same number of samples).
-
-    Args:
-        module: module to be parallelized
-        device_ids: CUDA devices (default: all devices)
-
-    Reference:
-        Hang Zhang, Kristin Dana, Jianping Shi, Zhongyue Zhang, Xiaogang Wang, Ambrish Tyagi,
-        Amit Agrawal. “Context Encoding for Semantic Segmentation.
-        *The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2018*
-
-    Example::
-
-        >>> net = encoding.nn.DataParallelModel(model, device_ids=[0, 1, 2])
-        >>> y = net(x)
-    """
-    def gather(self, outputs, output_device):
-        return outputs
-
-    def replicate(self, module, device_ids):
-        modules = super(DataParallelModel, self).replicate(module, device_ids)
-        return modules
-
-
-class DataParallelCriterion(DataParallel):
-    """
-    Calculate loss in multiple-GPUs, which balance the memory usage for
-    Semantic Segmentation.
-
-    The targets are splitted across the specified devices by chunking in
-    the batch dimension. Please use together with :class:`encoding.parallel.DataParallelModel`.
-
-    Reference:
-        Hang Zhang, Kristin Dana, Jianping Shi, Zhongyue Zhang, Xiaogang Wang, Ambrish Tyagi,
-        Amit Agrawal. “Context Encoding for Semantic Segmentation.
-        *The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) 2018*
-
-    Example::
-
-        >>> net = encoding.nn.DataParallelModel(model, device_ids=[0, 1, 2])
-        >>> criterion = encoding.nn.DataParallelCriterion(criterion, device_ids=[0, 1, 2])
-        >>> y = net(x)
-        >>> loss = criterion(y, target)
-    """
-    def forward(self, inputs, *targets, **kwargs):
-        # input should be already scatterd
-        # scattering the targets instead
-        if not self.device_ids:
-            return self.module(inputs, *targets, **kwargs)
-        targets, kwargs = self.scatter(targets, kwargs, self.device_ids)
-        if len(self.device_ids) == 1:
-            return self.module(inputs, *targets[0], **kwargs[0])
-        replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
-        outputs = _criterion_parallel_apply(replicas, inputs, targets, kwargs)
-        return Reduce.apply(*outputs) / len(outputs)
-
-
-def _criterion_parallel_apply(modules, inputs, targets, kwargs_tup=None, devices=None):
-    assert len(modules) == len(inputs)
-    assert len(targets) == len(inputs)
-    if kwargs_tup:
-        assert len(modules) == len(kwargs_tup)
-    else:
-        kwargs_tup = ({},) * len(modules)
-    if devices is not None:
-        assert len(modules) == len(devices)
-    else:
-        devices = [None] * len(modules)
-
-    lock = threading.Lock()
-    results = {}
-    if torch_ver != "0.3":
-        grad_enabled = torch.is_grad_enabled()
-
-    def _worker(i, module, input, target, kwargs, device=None):
-        if torch_ver != "0.3":
-            torch.set_grad_enabled(grad_enabled)
-        if device is None:
-            device = get_a_var(input).get_device()
-        try:
-            if not isinstance(input, tuple):
-                input = (input,)
-            with torch.cuda.device(device):
-                output = module(*(input + target), **kwargs)
-            with lock:
-                results[i] = output
-        except Exception as e:
-            with lock:
-                results[i] = e
-
-    if len(modules) > 1:
-        threads = [threading.Thread(target=_worker,
-                                    args=(i, module, input, target,
-                                          kwargs, device),)
-                   for i, (module, input, target, kwargs, device) in
-                   enumerate(zip(modules, inputs, targets, kwargs_tup, devices))]
-
-        for thread in threads:
-            thread.start()
-        for thread in threads:
-            thread.join()
-    else:
-        _worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0])
-
-    outputs = []
-    for i in range(len(inputs)):
-        output = results[i]
-        if isinstance(output, Exception):
-            raise output
-        outputs.append(output)
-    return outputs

idm_masking/humanparsing/tools/simple_extractor_dataset.py

@@ -1,89 +0,0 @@
-#!/usr/bin/env python
-# -*- encoding: utf-8 -*-
-
-"""
-@Author  :   Peike Li
-@Contact :   peike.li@yahoo.com
-@File    :   dataset.py
-@Time    :   8/30/19 9:12 PM
-@Desc    :   Dataset Definition
-@License :   This source code is licensed under the license found in the
-             LICENSE file in the root directory of this source tree.
-"""
-
-import os
-import pdb
-
-import cv2
-import numpy as np
-from PIL import Image
-from torch.utils import data
-from .transforms import get_affine_transform
-
-
-class SimpleFolderDataset(data.Dataset):
-    def __init__(self, root, input_size=[512, 512], transform=None):
-        self.root = root
-        self.input_size = input_size
-        self.transform = transform
-        self.aspect_ratio = input_size[1] * 1.0 / input_size[0]
-        self.input_size = np.asarray(input_size)
-        self.is_pil_image = False
-        if isinstance(root, Image.Image):
-            self.file_list = [root]
-            self.is_pil_image = True
-        elif os.path.isfile(root):
-            self.file_list = [os.path.basename(root)]
-            self.root = os.path.dirname(root)
-        else:
-            self.file_list = os.listdir(self.root)
-
-    def __len__(self):
-        return len(self.file_list)
-
-    def _box2cs(self, box):
-        x, y, w, h = box[:4]
-        return self._xywh2cs(x, y, w, h)
-
-    def _xywh2cs(self, x, y, w, h):
-        center = np.zeros((2), dtype=np.float32)
-        center[0] = x + w * 0.5
-        center[1] = y + h * 0.5
-        if w > self.aspect_ratio * h:
-            h = w * 1.0 / self.aspect_ratio
-        elif w < self.aspect_ratio * h:
-            w = h * self.aspect_ratio
-        scale = np.array([w, h], dtype=np.float32)
-        return center, scale
-
-    def __getitem__(self, index):
-        if self.is_pil_image:
-            img = np.asarray(self.file_list[index])[:, :, [2, 1, 0]]
-        else:
-            img_name = self.file_list[index]
-            img_path = os.path.join(self.root, img_name)
-            img = cv2.imread(img_path, cv2.IMREAD_COLOR)
-        h, w, _ = img.shape
-
-        # Get person center and scale
-        person_center, s = self._box2cs([0, 0, w - 1, h - 1])
-        r = 0
-        trans = get_affine_transform(person_center, s, r, self.input_size)
-        input = cv2.warpAffine(
-            img,
-            trans,
-            (int(self.input_size[1]), int(self.input_size[0])),
-            flags=cv2.INTER_LINEAR,
-            borderMode=cv2.BORDER_CONSTANT,
-            borderValue=(0, 0, 0))
-
-        input = self.transform(input)
-        meta = {
-            'center': person_center,
-            'height': h,
-            'width': w,
-            'scale': s,
-            'rotation': r
-        }
-
-        return input, meta

idm_masking/humanparsing/tools/transforms.py

@@ -1,167 +0,0 @@
-# ------------------------------------------------------------------------------
-# Copyright (c) Microsoft
-# Licensed under the MIT License.
-# Written by Bin Xiao (Bin.Xiao@microsoft.com)
-# ------------------------------------------------------------------------------
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import numpy as np
-import cv2
-import torch
-
-class BRG2Tensor_transform(object):
-    def __call__(self, pic):
-        img = torch.from_numpy(pic.transpose((2, 0, 1)))
-        if isinstance(img, torch.ByteTensor):
-            return img.float()
-        else:
-            return img
-
-class BGR2RGB_transform(object):
-    def __call__(self, tensor):
-        return tensor[[2,1,0],:,:]
-
-def flip_back(output_flipped, matched_parts):
-    '''
-    ouput_flipped: numpy.ndarray(batch_size, num_joints, height, width)
-    '''
-    assert output_flipped.ndim == 4,\
-        'output_flipped should be [batch_size, num_joints, height, width]'
-
-    output_flipped = output_flipped[:, :, :, ::-1]
-
-    for pair in matched_parts:
-        tmp = output_flipped[:, pair[0], :, :].copy()
-        output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :]
-        output_flipped[:, pair[1], :, :] = tmp
-
-    return output_flipped
-
-
-def fliplr_joints(joints, joints_vis, width, matched_parts):
-    """
-    flip coords
-    """
-    # Flip horizontal
-    joints[:, 0] = width - joints[:, 0] - 1
-
-    # Change left-right parts
-    for pair in matched_parts:
-        joints[pair[0], :], joints[pair[1], :] = \
-            joints[pair[1], :], joints[pair[0], :].copy()
-        joints_vis[pair[0], :], joints_vis[pair[1], :] = \
-            joints_vis[pair[1], :], joints_vis[pair[0], :].copy()
-
-    return joints*joints_vis, joints_vis
-
-
-def transform_preds(coords, center, scale, input_size):
-    target_coords = np.zeros(coords.shape)
-    trans = get_affine_transform(center, scale, 0, input_size, inv=1)
-    for p in range(coords.shape[0]):
-        target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
-    return target_coords
-
-def transform_parsing(pred, center, scale, width, height, input_size):
-
-    trans = get_affine_transform(center, scale, 0, input_size, inv=1)
-    target_pred = cv2.warpAffine(
-            pred,
-            trans,
-            (int(width), int(height)), #(int(width), int(height)),
-            flags=cv2.INTER_NEAREST,
-            borderMode=cv2.BORDER_CONSTANT,
-            borderValue=(0))
-
-    return target_pred
-
-def transform_logits(logits, center, scale, width, height, input_size):
-
-    trans = get_affine_transform(center, scale, 0, input_size, inv=1)
-    channel = logits.shape[2]
-    target_logits = []
-    for i in range(channel):
-        target_logit = cv2.warpAffine(
-            logits[:,:,i],
-            trans,
-            (int(width), int(height)), #(int(width), int(height)),
-            flags=cv2.INTER_LINEAR,
-            borderMode=cv2.BORDER_CONSTANT,
-            borderValue=(0))
-        target_logits.append(target_logit)
-    target_logits = np.stack(target_logits,axis=2)
-
-    return target_logits
-
-
-def get_affine_transform(center,
-                         scale,
-                         rot,
-                         output_size,
-                         shift=np.array([0, 0], dtype=np.float32),
-                         inv=0):
-    if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
-        print(scale)
-        scale = np.array([scale, scale])
-
-    scale_tmp = scale
-
-    src_w = scale_tmp[0]
-    dst_w = output_size[1]
-    dst_h = output_size[0]
-
-    rot_rad = np.pi * rot / 180
-    src_dir = get_dir([0, src_w * -0.5], rot_rad)
-    dst_dir = np.array([0, (dst_w-1) * -0.5], np.float32)
-
-    src = np.zeros((3, 2), dtype=np.float32)
-    dst = np.zeros((3, 2), dtype=np.float32)
-    src[0, :] = center + scale_tmp * shift
-    src[1, :] = center + src_dir + scale_tmp * shift
-    dst[0, :] = [(dst_w-1) * 0.5, (dst_h-1) * 0.5]
-    dst[1, :] = np.array([(dst_w-1) * 0.5, (dst_h-1) * 0.5]) + dst_dir
-
-    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
-    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
-
-    if inv:
-        trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
-    else:
-        trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
-
-    return trans
-
-
-def affine_transform(pt, t):
-    new_pt = np.array([pt[0], pt[1], 1.]).T
-    new_pt = np.dot(t, new_pt)
-    return new_pt[:2]
-
-
-def get_3rd_point(a, b):
-    direct = a - b
-    return b + np.array([-direct[1], direct[0]], dtype=np.float32)
-
-
-def get_dir(src_point, rot_rad):
-    sn, cs = np.sin(rot_rad), np.cos(rot_rad)
-
-    src_result = [0, 0]
-    src_result[0] = src_point[0] * cs - src_point[1] * sn
-    src_result[1] = src_point[0] * sn + src_point[1] * cs
-
-    return src_result
-
-
-def crop(img, center, scale, output_size, rot=0):
-    trans = get_affine_transform(center, scale, rot, output_size)
-
-    dst_img = cv2.warpAffine(img,
-                             trans,
-                             (int(output_size[1]), int(output_size[0])),
-                             flags=cv2.INTER_LINEAR)
-
-    return dst_img

idm_masking/openpose/annotator/openpose/LICENSE

@@ -1,108 +0,0 @@
-OPENPOSE: MULTIPERSON KEYPOINT DETECTION
-SOFTWARE LICENSE AGREEMENT
-ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY
-
-BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS LICENSE AGREEMENT.  IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR DOWNLOAD THE SOFTWARE.
-
-This is a license agreement ("Agreement") between your academic institution or non-profit organization or self (called "Licensee" or "You" in this Agreement) and Carnegie Mellon University (called "Licensor" in this Agreement).  All rights not specifically granted to you in this Agreement are reserved for Licensor. 
-
-RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: 
-Licensor retains exclusive ownership of any copy of the Software (as defined below) licensed under this Agreement and hereby grants to Licensee a personal, non-exclusive, 
-non-transferable license to use the Software for noncommercial research purposes, without the right to sublicense, pursuant to the terms and conditions of this Agreement.  As used in this Agreement, the term "Software" means (i) the actual copy of all or any portion of code for program routines made accessible to Licensee by Licensor pursuant to this Agreement, inclusive of backups, updates, and/or merged copies permitted hereunder or subsequently supplied by Licensor,  including all or any file structures, programming instructions, user interfaces and screen formats and sequences as well as any and all documentation and instructions related to it, and (ii) all or any derivatives and/or modifications created or made by You to any of the items specified in (i).
-
-CONFIDENTIALITY: Licensee acknowledges that the Software is proprietary to Licensor, and as such, Licensee agrees to receive all such materials in confidence and use the Software only in accordance with the terms of this Agreement.  Licensee agrees to use reasonable effort to protect the Software from unauthorized use, reproduction, distribution, or publication.
-
-COPYRIGHT: The Software is owned by Licensor and is protected by United 
-States copyright laws and applicable international treaties and/or conventions.
-
-PERMITTED USES:  The Software may be used for your own noncommercial internal research purposes. You understand and agree that Licensor is not obligated to implement any suggestions and/or feedback you might provide regarding the Software, but to the extent Licensor does so, you are not entitled to any compensation related thereto.
-
-DERIVATIVES: You may create derivatives of or make modifications to the Software, however, You agree that all and any such derivatives and modifications will be owned by Licensor and become a part of the Software licensed to You under this Agreement.  You may only use such derivatives and modifications for your own noncommercial internal research purposes, and you may not otherwise use, distribute or copy such derivatives and modifications in violation of this Agreement.
-
-BACKUPS:  If Licensee is an organization, it may make that number of copies of the Software necessary for internal noncommercial use at a single site within its organization provided that all information appearing in or on the original labels, including the copyright and trademark notices are copied onto the labels of the copies.
-
-USES NOT PERMITTED:  You may not distribute, copy or use the Software except as explicitly permitted herein. Licensee has not been granted any trademark license as part of this Agreement and may not use the name or mark “OpenPose", "Carnegie Mellon" or any renditions thereof without the prior written permission of Licensor.
-
-You may not sell, rent, lease, sublicense, lend, time-share or transfer, in whole or in part, or provide third parties access to prior or present versions (or any parts thereof) of the Software.
-
-ASSIGNMENT: You may not assign this Agreement or your rights hereunder without the prior written consent of Licensor. Any attempted assignment without such consent shall be null and void.
-
-TERM: The term of the license granted by this Agreement is from Licensee's acceptance of this Agreement by downloading the Software or by using the Software until terminated as provided below.
-
-The Agreement automatically terminates without notice if you fail to comply with any provision of this Agreement.  Licensee may terminate this Agreement by ceasing using the Software.  Upon any termination of this Agreement, Licensee will delete any and all copies of the Software. You agree that all provisions which operate to protect the proprietary rights of Licensor shall remain in force should breach occur and that the obligation of confidentiality described in this Agreement is binding in perpetuity and, as such, survives the term of the Agreement.
-
-FEE: Provided Licensee abides completely by the terms and conditions of this Agreement, there is no fee due to Licensor for Licensee's use of the Software in accordance with this Agreement.
-
-DISCLAIMER OF WARRANTIES:  THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON-INFRINGEMENT.  LICENSEE BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND RELATED MATERIALS.
-
-SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is provided as part of this Agreement.  
-
-EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent permitted under applicable law, Licensor shall not be liable for direct, indirect, special, incidental, or consequential damages or lost profits related to Licensee's use of and/or inability to use the Software, even if Licensor is advised of the possibility of such damage.
-
-EXPORT REGULATION: Licensee agrees to comply with any and all applicable 
-U.S. export control laws, regulations, and/or other laws related to embargoes and sanction programs administered by the Office of Foreign Assets Control.
-
-SEVERABILITY: If any provision(s) of this Agreement shall be held to be invalid, illegal, or unenforceable by a court or other tribunal of competent jurisdiction, the validity, legality and enforceability of the remaining provisions shall not in any way be affected or impaired thereby.
-
-NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right or remedy under this Agreement shall be construed as a waiver of any future or other exercise of such right or remedy by Licensor.
-
-GOVERNING LAW: This Agreement shall be construed and enforced in accordance with the laws of the Commonwealth of Pennsylvania without reference to conflict of laws principles.  You consent to the personal jurisdiction of the courts of this County and waive their rights to venue outside of Allegheny County, Pennsylvania.
-
-ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and entire agreement between Licensee and Licensor as to the matter set forth herein and supersedes any previous agreements, understandings, and arrangements between the parties relating hereto.
-
-
-
-************************************************************************
-
-THIRD-PARTY SOFTWARE NOTICES AND INFORMATION
-
-This project incorporates material from the project(s) listed below (collectively, "Third Party Code").  This Third Party Code is licensed to you under their original license terms set forth below.  We reserves all other rights not expressly granted, whether by implication, estoppel or otherwise.
- 
-1.	Caffe, version 1.0.0, (https://github.com/BVLC/caffe/)
-
-COPYRIGHT
-
-All contributions by the University of California:
-Copyright (c) 2014-2017 The Regents of the University of California (Regents)
-All rights reserved.
-
-All other contributions:
-Copyright (c) 2014-2017, the respective contributors
-All rights reserved.
-
-Caffe uses a shared copyright model: each contributor holds copyright over
-their contributions to Caffe. The project versioning records all such
-contribution and copyright details. If a contributor wants to further mark
-their specific copyright on a particular contribution, they should indicate
-their copyright solely in the commit message of the change when it is
-committed.
-
-LICENSE
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met: 
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer. 
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution. 
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
-ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
-WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
-ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
-(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
-LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
-ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
-(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
-SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-
-CONTRIBUTION AGREEMENT
-
-By contributing to the BVLC/caffe repository through pull-request, comment,
-or otherwise, the contributor releases their content to the
-license and copyright terms herein.
-
-************END OF THIRD-PARTY SOFTWARE NOTICES AND INFORMATION**********

idm_masking/openpose/annotator/openpose/__init__.py

@@ -1,102 +0,0 @@
-# Openpose
-# Original from CMU https://github.com/CMU-Perceptual-Computing-Lab/openpose
-# 2nd Edited by https://github.com/Hzzone/pytorch-openpose
-# 3rd Edited by ControlNet
-# 4th Edited by ControlNet (added face and correct hands)
-
-import os
-import pdb
-
-os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
-
-import torch
-import numpy as np
-from . import util
-from .body import Body
-from .hand import Hand
-from .face import Face
-from annotator.util import annotator_ckpts_path
-
-body_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/body_pose_model.pth"
-hand_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/hand_pose_model.pth"
-face_model_path = "https://huggingface.co/lllyasviel/Annotators/resolve/main/facenet.pth"
-
-
-def draw_pose(pose, H, W, draw_body=True, draw_hand=True, draw_face=True):
-    bodies = pose['bodies']
-    faces = pose['faces']
-    hands = pose['hands']
-    candidate = bodies['candidate']
-    subset = bodies['subset']
-    canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
-
-    if draw_body:
-        canvas = util.draw_bodypose(canvas, candidate, subset)
-
-    if draw_hand:
-        canvas = util.draw_handpose(canvas, hands)
-
-    if draw_face:
-        canvas = util.draw_facepose(canvas, faces)
-
-    return canvas
-
-
-class OpenposeDetector:
-    def __init__(self):
-        body_modelpath = os.path.join(annotator_ckpts_path, "body_pose_model.pth")
-        # hand_modelpath = os.path.join(annotator_ckpts_path, "hand_pose_model.pth")
-        # face_modelpath = os.path.join(annotator_ckpts_path, "facenet.pth")
-
-        if not os.path.exists(body_modelpath):
-            from basicsr.utils.download_util import load_file_from_url
-            load_file_from_url(body_model_path, model_dir=annotator_ckpts_path)
-
-        # if not os.path.exists(hand_modelpath):
-        #     from basicsr.utils.download_util import load_file_from_url
-        #     load_file_from_url(hand_model_path, model_dir=annotator_ckpts_path)
-
-        # if not os.path.exists(face_modelpath):
-        #     from basicsr.utils.download_util import load_file_from_url
-        #     load_file_from_url(face_model_path, model_dir=annotator_ckpts_path)
-
-        self.body_estimation = Body(body_modelpath)
-        # self.hand_estimation = Hand(hand_modelpath)
-        # self.face_estimation = Face(face_modelpath)
-
-
-    def __call__(self, oriImg, hand_and_face=False, return_is_index=False):
-        oriImg = oriImg[:, :, ::-1].copy()
-        H, W, C = oriImg.shape
-        with torch.no_grad():
-            candidate, subset = self.body_estimation(oriImg)
-            hands = []
-            faces = []
-            if hand_and_face:
-                # Hand
-                hands_list = util.handDetect(candidate, subset, oriImg)
-                for x, y, w, is_left in hands_list:
-                    peaks = self.hand_estimation(oriImg[y:y + w, x:x + w, :]).astype(np.float32)
-                    if peaks.ndim == 2 and peaks.shape[1] == 2:
-                        peaks[:, 0] = np.where(peaks[:, 0] < 1e-6, -1, peaks[:, 0] + x) / float(W)
-                        peaks[:, 1] = np.where(peaks[:, 1] < 1e-6, -1, peaks[:, 1] + y) / float(H)
-                        hands.append(peaks.tolist())
-                # Face
-                faces_list = util.faceDetect(candidate, subset, oriImg)
-                for x, y, w in faces_list:
-                    heatmaps = self.face_estimation(oriImg[y:y + w, x:x + w, :])
-                    peaks = self.face_estimation.compute_peaks_from_heatmaps(heatmaps).astype(np.float32)
-                    if peaks.ndim == 2 and peaks.shape[1] == 2:
-                        peaks[:, 0] = np.where(peaks[:, 0] < 1e-6, -1, peaks[:, 0] + x) / float(W)
-                        peaks[:, 1] = np.where(peaks[:, 1] < 1e-6, -1, peaks[:, 1] + y) / float(H)
-                        faces.append(peaks.tolist())
-            if candidate.ndim == 2 and candidate.shape[1] == 4:
-                candidate = candidate[:, :2]
-                candidate[:, 0] /= float(W)
-                candidate[:, 1] /= float(H)
-            bodies = dict(candidate=candidate.tolist(), subset=subset.tolist())
-            pose = dict(bodies=bodies, hands=hands, faces=faces)
-            if return_is_index:
-                return pose
-            else:
-                return pose, draw_pose(pose, H, W)

idm_masking/openpose/annotator/openpose/body.py

@@ -1,211 +0,0 @@
-from pathlib import Path
-import sys
-PROJECT_ROOT = Path(__file__).absolute().parents[3].absolute()
-# print(PROJECT_ROOT)
-
-import cv2
-import numpy as np
-import math
-import time
-from scipy.ndimage.filters import gaussian_filter
-import matplotlib.pyplot as plt
-import matplotlib
-import torch
-from torchvision import transforms
-
-from . import util
-from .model import bodypose_model
-
-
-class Body(object):
-    def __init__(self, model_path):
-        self.model = bodypose_model()
-        if torch.cuda.is_available():
-            self.model = self.model.cuda()
-            # print('cuda')
-        model_dict = util.transfer(self.model, torch.load(model_path))
-        self.model.load_state_dict(model_dict)
-        self.model.eval()
-
-
-    def __call__(self, oriImg):
-        # scale_search = [0.5, 1.0, 1.5, 2.0]
-        scale_search = [0.5]
-        boxsize = 368
-        stride = 8
-        padValue = 128
-        thre1 = 0.1
-        thre2 = 0.05
-        multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
-        heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
-        paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
-
-        for m in range(len(multiplier)):
-            scale = multiplier[m]
-            imageToTest = util.smart_resize_k(oriImg, fx=scale, fy=scale)
-            imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
-            im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
-            im = np.ascontiguousarray(im)
-
-            data = torch.from_numpy(im).float()
-            if torch.cuda.is_available():
-                data = data.cuda()
-            with torch.no_grad():
-                Mconv7_stage6_L1, Mconv7_stage6_L2 = self.model(data)
-
-            Mconv7_stage6_L1 = Mconv7_stage6_L1.cpu().numpy()
-            Mconv7_stage6_L2 = Mconv7_stage6_L2.cpu().numpy()
-
-            heatmap = np.transpose(np.squeeze(Mconv7_stage6_L2), (1, 2, 0))
-            heatmap = util.smart_resize_k(heatmap, fx=stride, fy=stride)
-            heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
-            heatmap = util.smart_resize(heatmap, (oriImg.shape[0], oriImg.shape[1]))
-
-            paf = np.transpose(np.squeeze(Mconv7_stage6_L1), (1, 2, 0))
-            paf = util.smart_resize_k(paf, fx=stride, fy=stride)
-            paf = paf[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
-            paf = util.smart_resize(paf, (oriImg.shape[0], oriImg.shape[1]))
-
-            heatmap_avg += heatmap_avg + heatmap / len(multiplier)
-            paf_avg += + paf / len(multiplier)
-
-        all_peaks = []
-        peak_counter = 0
-
-        for part in range(18):
-            map_ori = heatmap_avg[:, :, part]
-            one_heatmap = gaussian_filter(map_ori, sigma=3)
-
-            map_left = np.zeros(one_heatmap.shape)
-            map_left[1:, :] = one_heatmap[:-1, :]
-            map_right = np.zeros(one_heatmap.shape)
-            map_right[:-1, :] = one_heatmap[1:, :]
-            map_up = np.zeros(one_heatmap.shape)
-            map_up[:, 1:] = one_heatmap[:, :-1]
-            map_down = np.zeros(one_heatmap.shape)
-            map_down[:, :-1] = one_heatmap[:, 1:]
-
-            peaks_binary = np.logical_and.reduce(
-                (one_heatmap >= map_left, one_heatmap >= map_right, one_heatmap >= map_up, one_heatmap >= map_down,
-                 one_heatmap > thre1))
-            peaks = list(zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]))  # note reverse
-            peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks]
-            peak_id = range(peak_counter, peak_counter + len(peaks))
-            peaks_with_score_and_id = [peaks_with_score[i] + (peak_id[i],) for i in range(len(peak_id))]
-
-            all_peaks.append(peaks_with_score_and_id)
-            peak_counter += len(peaks)
-
-        # find connection in the specified sequence, center 29 is in the position 15
-        limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
-                   [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
-                   [1, 16], [16, 18], [3, 17], [6, 18]]
-        # the middle joints heatmap correpondence
-        mapIdx = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44], [19, 20], [21, 22], \
-                  [23, 24], [25, 26], [27, 28], [29, 30], [47, 48], [49, 50], [53, 54], [51, 52], \
-                  [55, 56], [37, 38], [45, 46]]
-
-        connection_all = []
-        special_k = []
-        mid_num = 10
-
-        for k in range(len(mapIdx)):
-            score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
-            candA = all_peaks[limbSeq[k][0] - 1]
-            candB = all_peaks[limbSeq[k][1] - 1]
-            nA = len(candA)
-            nB = len(candB)
-            indexA, indexB = limbSeq[k]
-            if (nA != 0 and nB != 0):
-                connection_candidate = []
-                for i in range(nA):
-                    for j in range(nB):
-                        vec = np.subtract(candB[j][:2], candA[i][:2])
-                        norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
-                        norm = max(0.001, norm)
-                        vec = np.divide(vec, norm)
-
-                        startend = list(zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
-                                            np.linspace(candA[i][1], candB[j][1], num=mid_num)))
-
-                        vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
-                                          for I in range(len(startend))])
-                        vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
-                                          for I in range(len(startend))])
-
-                        score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
-                        score_with_dist_prior = sum(score_midpts) / len(score_midpts) + min(
-                            0.5 * oriImg.shape[0] / norm - 1, 0)
-                        criterion1 = len(np.nonzero(score_midpts > thre2)[0]) > 0.8 * len(score_midpts)
-                        criterion2 = score_with_dist_prior > 0
-                        if criterion1 and criterion2:
-                            connection_candidate.append(
-                                [i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2]])
-
-                connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
-                connection = np.zeros((0, 5))
-                for c in range(len(connection_candidate)):
-                    i, j, s = connection_candidate[c][0:3]
-                    if (i not in connection[:, 3] and j not in connection[:, 4]):
-                        connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
-                        if (len(connection) >= min(nA, nB)):
-                            break
-
-                connection_all.append(connection)
-            else:
-                special_k.append(k)
-                connection_all.append([])
-
-        subset = -1 * np.ones((0, 20))
-        candidate = np.array([item for sublist in all_peaks for item in sublist])
-
-        for k in range(len(mapIdx)):
-            if k not in special_k:
-                partAs = connection_all[k][:, 0]
-                partBs = connection_all[k][:, 1]
-                indexA, indexB = np.array(limbSeq[k]) - 1
-
-                for i in range(len(connection_all[k])):
-                    found = 0
-                    subset_idx = [-1, -1]
-                    for j in range(len(subset)):
-                        if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
-                            subset_idx[found] = j
-                            found += 1
-
-                    if found == 1:
-                        j = subset_idx[0]
-                        if subset[j][indexB] != partBs[i]:
-                            subset[j][indexB] = partBs[i]
-                            subset[j][-1] += 1
-                            subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
-                    elif found == 2:
-                        j1, j2 = subset_idx
-                        membership = ((subset[j1] >= 0).astype(int) + (subset[j2] >= 0).astype(int))[:-2]
-                        if len(np.nonzero(membership == 2)[0]) == 0:
-                            subset[j1][:-2] += (subset[j2][:-2] + 1)
-                            subset[j1][-2:] += subset[j2][-2:]
-                            subset[j1][-2] += connection_all[k][i][2]
-                            subset = np.delete(subset, j2, 0)
-                        else:
-                            subset[j1][indexB] = partBs[i]
-                            subset[j1][-1] += 1
-                            subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
-
-
-                    elif not found and k < 17:
-                        row = -1 * np.ones(20)
-                        row[indexA] = partAs[i]
-                        row[indexB] = partBs[i]
-                        row[-1] = 2
-                        row[-2] = sum(candidate[connection_all[k][i, :2].astype(int), 2]) + connection_all[k][i][2]
-                        subset = np.vstack([subset, row])
-
-        deleteIdx = []
-        for i in range(len(subset)):
-            if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
-                deleteIdx.append(i)
-        subset = np.delete(subset, deleteIdx, axis=0)
-
-
-        return candidate, subset

idm_masking/openpose/annotator/openpose/face.py

@@ -1,368 +0,0 @@
-import logging
-import numpy as np
-from torchvision.transforms import ToTensor, ToPILImage
-import torch
-import torch.nn.functional as F
-import cv2
-
-from . import util
-from torch.nn import Conv2d, Module, ReLU, MaxPool2d, init
-
-
-class FaceNet(Module):
-    """Model the cascading heatmaps. """
-    def __init__(self):
-        super(FaceNet, self).__init__()
-        # cnn to make feature map
-        self.relu = ReLU()
-        self.max_pooling_2d = MaxPool2d(kernel_size=2, stride=2)
-        self.conv1_1 = Conv2d(in_channels=3, out_channels=64,
-                              kernel_size=3, stride=1, padding=1)
-        self.conv1_2 = Conv2d(
-            in_channels=64, out_channels=64, kernel_size=3, stride=1,
-            padding=1)
-        self.conv2_1 = Conv2d(
-            in_channels=64, out_channels=128, kernel_size=3, stride=1,
-            padding=1)
-        self.conv2_2 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=3, stride=1,
-            padding=1)
-        self.conv3_1 = Conv2d(
-            in_channels=128, out_channels=256, kernel_size=3, stride=1,
-            padding=1)
-        self.conv3_2 = Conv2d(
-            in_channels=256, out_channels=256, kernel_size=3, stride=1,
-            padding=1)
-        self.conv3_3 = Conv2d(
-            in_channels=256, out_channels=256, kernel_size=3, stride=1,
-            padding=1)
-        self.conv3_4 = Conv2d(
-            in_channels=256, out_channels=256, kernel_size=3, stride=1,
-            padding=1)
-        self.conv4_1 = Conv2d(
-            in_channels=256, out_channels=512, kernel_size=3, stride=1,
-            padding=1)
-        self.conv4_2 = Conv2d(
-            in_channels=512, out_channels=512, kernel_size=3, stride=1,
-            padding=1)
-        self.conv4_3 = Conv2d(
-            in_channels=512, out_channels=512, kernel_size=3, stride=1,
-            padding=1)
-        self.conv4_4 = Conv2d(
-            in_channels=512, out_channels=512, kernel_size=3, stride=1,
-            padding=1)
-        self.conv5_1 = Conv2d(
-            in_channels=512, out_channels=512, kernel_size=3, stride=1,
-            padding=1)
-        self.conv5_2 = Conv2d(
-            in_channels=512, out_channels=512, kernel_size=3, stride=1,
-            padding=1)
-        self.conv5_3_CPM = Conv2d(
-            in_channels=512, out_channels=128, kernel_size=3, stride=1,
-            padding=1)
-
-        # stage1
-        self.conv6_1_CPM = Conv2d(
-            in_channels=128, out_channels=512, kernel_size=1, stride=1,
-            padding=0)
-        self.conv6_2_CPM = Conv2d(
-            in_channels=512, out_channels=71, kernel_size=1, stride=1,
-            padding=0)
-
-        # stage2
-        self.Mconv1_stage2 = Conv2d(
-            in_channels=199, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv2_stage2 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv3_stage2 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv4_stage2 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv5_stage2 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv6_stage2 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=1, stride=1,
-            padding=0)
-        self.Mconv7_stage2 = Conv2d(
-            in_channels=128, out_channels=71, kernel_size=1, stride=1,
-            padding=0)
-
-        # stage3
-        self.Mconv1_stage3 = Conv2d(
-            in_channels=199, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv2_stage3 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv3_stage3 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv4_stage3 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv5_stage3 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv6_stage3 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=1, stride=1,
-            padding=0)
-        self.Mconv7_stage3 = Conv2d(
-            in_channels=128, out_channels=71, kernel_size=1, stride=1,
-            padding=0)
-
-        # stage4
-        self.Mconv1_stage4 = Conv2d(
-            in_channels=199, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv2_stage4 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv3_stage4 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv4_stage4 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv5_stage4 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv6_stage4 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=1, stride=1,
-            padding=0)
-        self.Mconv7_stage4 = Conv2d(
-            in_channels=128, out_channels=71, kernel_size=1, stride=1,
-            padding=0)
-
-        # stage5
-        self.Mconv1_stage5 = Conv2d(
-            in_channels=199, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv2_stage5 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv3_stage5 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv4_stage5 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv5_stage5 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv6_stage5 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=1, stride=1,
-            padding=0)
-        self.Mconv7_stage5 = Conv2d(
-            in_channels=128, out_channels=71, kernel_size=1, stride=1,
-            padding=0)
-
-        # stage6
-        self.Mconv1_stage6 = Conv2d(
-            in_channels=199, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv2_stage6 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv3_stage6 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv4_stage6 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv5_stage6 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=7, stride=1,
-            padding=3)
-        self.Mconv6_stage6 = Conv2d(
-            in_channels=128, out_channels=128, kernel_size=1, stride=1,
-            padding=0)
-        self.Mconv7_stage6 = Conv2d(
-            in_channels=128, out_channels=71, kernel_size=1, stride=1,
-            padding=0)
-
-        for m in self.modules():
-            if isinstance(m, Conv2d):
-                init.constant_(m.bias, 0)
-
-    def forward(self, x):
-        """Return a list of heatmaps."""
-        heatmaps = []
-
-        h = self.relu(self.conv1_1(x))
-        h = self.relu(self.conv1_2(h))
-        h = self.max_pooling_2d(h)
-        h = self.relu(self.conv2_1(h))
-        h = self.relu(self.conv2_2(h))
-        h = self.max_pooling_2d(h)
-        h = self.relu(self.conv3_1(h))
-        h = self.relu(self.conv3_2(h))
-        h = self.relu(self.conv3_3(h))
-        h = self.relu(self.conv3_4(h))
-        h = self.max_pooling_2d(h)
-        h = self.relu(self.conv4_1(h))
-        h = self.relu(self.conv4_2(h))
-        h = self.relu(self.conv4_3(h))
-        h = self.relu(self.conv4_4(h))
-        h = self.relu(self.conv5_1(h))
-        h = self.relu(self.conv5_2(h))
-        h = self.relu(self.conv5_3_CPM(h))
-        feature_map = h
-
-        # stage1
-        h = self.relu(self.conv6_1_CPM(h))
-        h = self.conv6_2_CPM(h)
-        heatmaps.append(h)
-
-        # stage2
-        h = torch.cat([h, feature_map], dim=1)  # channel concat
-        h = self.relu(self.Mconv1_stage2(h))
-        h = self.relu(self.Mconv2_stage2(h))
-        h = self.relu(self.Mconv3_stage2(h))
-        h = self.relu(self.Mconv4_stage2(h))
-        h = self.relu(self.Mconv5_stage2(h))
-        h = self.relu(self.Mconv6_stage2(h))
-        h = self.Mconv7_stage2(h)
-        heatmaps.append(h)
-
-        # stage3
-        h = torch.cat([h, feature_map], dim=1)  # channel concat
-        h = self.relu(self.Mconv1_stage3(h))
-        h = self.relu(self.Mconv2_stage3(h))
-        h = self.relu(self.Mconv3_stage3(h))
-        h = self.relu(self.Mconv4_stage3(h))
-        h = self.relu(self.Mconv5_stage3(h))
-        h = self.relu(self.Mconv6_stage3(h))
-        h = self.Mconv7_stage3(h)
-        heatmaps.append(h)
-
-        # stage4
-        h = torch.cat([h, feature_map], dim=1)  # channel concat
-        h = self.relu(self.Mconv1_stage4(h))
-        h = self.relu(self.Mconv2_stage4(h))
-        h = self.relu(self.Mconv3_stage4(h))
-        h = self.relu(self.Mconv4_stage4(h))
-        h = self.relu(self.Mconv5_stage4(h))
-        h = self.relu(self.Mconv6_stage4(h))
-        h = self.Mconv7_stage4(h)
-        heatmaps.append(h)
-
-        # stage5
-        h = torch.cat([h, feature_map], dim=1)  # channel concat
-        h = self.relu(self.Mconv1_stage5(h))
-        h = self.relu(self.Mconv2_stage5(h))
-        h = self.relu(self.Mconv3_stage5(h))
-        h = self.relu(self.Mconv4_stage5(h))
-        h = self.relu(self.Mconv5_stage5(h))
-        h = self.relu(self.Mconv6_stage5(h))
-        h = self.Mconv7_stage5(h)
-        heatmaps.append(h)
-
-        # stage6
-        h = torch.cat([h, feature_map], dim=1)  # channel concat
-        h = self.relu(self.Mconv1_stage6(h))
-        h = self.relu(self.Mconv2_stage6(h))
-        h = self.relu(self.Mconv3_stage6(h))
-        h = self.relu(self.Mconv4_stage6(h))
-        h = self.relu(self.Mconv5_stage6(h))
-        h = self.relu(self.Mconv6_stage6(h))
-        h = self.Mconv7_stage6(h)
-        heatmaps.append(h)
-
-        return heatmaps
-
-
-LOG = logging.getLogger(__name__)
-TOTEN = ToTensor()
-TOPIL = ToPILImage()
-
-
-params = {
-    'gaussian_sigma': 2.5,
-    'inference_img_size': 736,  # 368, 736, 1312
-    'heatmap_peak_thresh': 0.1,
-    'crop_scale': 1.5,
-    'line_indices': [
-        [0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 6],
-        [6, 7], [7, 8], [8, 9], [9, 10], [10, 11], [11, 12], [12, 13],
-        [13, 14], [14, 15], [15, 16],
-        [17, 18], [18, 19], [19, 20], [20, 21],
-        [22, 23], [23, 24], [24, 25], [25, 26],
-        [27, 28], [28, 29], [29, 30],
-        [31, 32], [32, 33], [33, 34], [34, 35],
-        [36, 37], [37, 38], [38, 39], [39, 40], [40, 41], [41, 36],
-        [42, 43], [43, 44], [44, 45], [45, 46], [46, 47], [47, 42],
-        [48, 49], [49, 50], [50, 51], [51, 52], [52, 53], [53, 54],
-        [54, 55], [55, 56], [56, 57], [57, 58], [58, 59], [59, 48],
-        [60, 61], [61, 62], [62, 63], [63, 64], [64, 65], [65, 66],
-        [66, 67], [67, 60]
-    ],
-}
-
-
-class Face(object):
-    """
-    The OpenPose face landmark detector model.
-
-    Args:
-        inference_size: set the size of the inference image size, suggested:
-            368, 736, 1312, default 736
-        gaussian_sigma: blur the heatmaps, default 2.5
-        heatmap_peak_thresh: return landmark if over threshold, default 0.1
-
-    """
-    def __init__(self, face_model_path,
-                 inference_size=None,
-                 gaussian_sigma=None,
-                 heatmap_peak_thresh=None):
-        self.inference_size = inference_size or params["inference_img_size"]
-        self.sigma = gaussian_sigma or params['gaussian_sigma']
-        self.threshold = heatmap_peak_thresh or params["heatmap_peak_thresh"]
-        self.model = FaceNet()
-        self.model.load_state_dict(torch.load(face_model_path))
-        if torch.cuda.is_available():
-            self.model = self.model.cuda()
-            print('cuda')
-        self.model.eval()
-
-    def __call__(self, face_img):
-        H, W, C = face_img.shape
-
-        w_size = 384
-        x_data = torch.from_numpy(util.smart_resize(face_img, (w_size, w_size))).permute([2, 0, 1]) / 256.0 - 0.5
-
-        if torch.cuda.is_available():
-            x_data = x_data.cuda()
-
-        with torch.no_grad():
-            hs = self.model(x_data[None, ...])
-
-            # output_path = "/home/aigc/ProjectVTON/WebDemo/onnx_models/face_estimation.onnx"
-            # torch.onnx.export(self.model, x_data[None, ...], output_path, export_params=True,
-            #                   opset_version=11,
-            #                   do_constant_folding=True)
-            heatmaps = F.interpolate(
-                hs[-1],
-                (H, W),
-                mode='bilinear', align_corners=True).cpu().numpy()[0]
-        return heatmaps
-
-    def compute_peaks_from_heatmaps(self, heatmaps):
-        all_peaks = []
-        for part in range(heatmaps.shape[0]):
-            map_ori = heatmaps[part].copy()
-            binary = np.ascontiguousarray(map_ori > 0.05, dtype=np.uint8)
-
-            if np.sum(binary) == 0:
-                continue
-
-            positions = np.where(binary > 0.5)
-            intensities = map_ori[positions]
-            mi = np.argmax(intensities)
-            y, x = positions[0][mi], positions[1][mi]
-            all_peaks.append([x, y])
-
-        return np.array(all_peaks)

idm_masking/openpose/annotator/openpose/hand.py

@@ -1,98 +0,0 @@
-import cv2
-import json
-import numpy as np
-import math
-import time
-from scipy.ndimage.filters import gaussian_filter
-import matplotlib.pyplot as plt
-import matplotlib
-import torch
-from skimage.measure import label
-
-from .model import handpose_model
-from . import util
-
-class Hand(object):
-    def __init__(self, model_path):
-        self.model = handpose_model()
-        if torch.cuda.is_available():
-            self.model = self.model.cuda()
-            print('cuda')
-        model_dict = util.transfer(self.model, torch.load(model_path))
-        self.model.load_state_dict(model_dict)
-        self.model.eval()
-
-    def __call__(self, oriImgRaw):
-        scale_search = [0.5, 1.0, 1.5, 2.0]
-        # scale_search = [0.5]
-        boxsize = 368
-        stride = 8
-        padValue = 128
-        thre = 0.05
-        multiplier = [x * boxsize for x in scale_search]
-
-        wsize = 128
-        heatmap_avg = np.zeros((wsize, wsize, 22))
-
-        Hr, Wr, Cr = oriImgRaw.shape
-
-        oriImg = cv2.GaussianBlur(oriImgRaw, (0, 0), 0.8)
-
-        for m in range(len(multiplier)):
-            scale = multiplier[m]
-            imageToTest = util.smart_resize(oriImg, (scale, scale))
-
-            imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
-            im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
-            im = np.ascontiguousarray(im)
-
-            data = torch.from_numpy(im).float()
-            if torch.cuda.is_available():
-                data = data.cuda()
-
-            with torch.no_grad():
-                output = self.model(data).cpu().numpy()
-
-                # output_path = "/home/aigc/ProjectVTON/WebDemo/onnx_models/hand_estimation.onnx"
-                # torch.onnx.export(self.model, data, output_path, export_params=True,
-                #                   opset_version=11,
-                #                   do_constant_folding=True)
-
-            # extract outputs, resize, and remove padding
-            heatmap = np.transpose(np.squeeze(output), (1, 2, 0))  # output 1 is heatmaps
-            heatmap = util.smart_resize_k(heatmap, fx=stride, fy=stride)
-            heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
-            heatmap = util.smart_resize(heatmap, (wsize, wsize))
-
-            heatmap_avg += heatmap / len(multiplier)
-
-        all_peaks = []
-        for part in range(21):
-            map_ori = heatmap_avg[:, :, part]
-            one_heatmap = gaussian_filter(map_ori, sigma=3)
-            binary = np.ascontiguousarray(one_heatmap > thre, dtype=np.uint8)
-
-            if np.sum(binary) == 0:
-                all_peaks.append([0, 0])
-                continue
-            label_img, label_numbers = label(binary, return_num=True, connectivity=binary.ndim)
-            max_index = np.argmax([np.sum(map_ori[label_img == i]) for i in range(1, label_numbers + 1)]) + 1
-            label_img[label_img != max_index] = 0
-            map_ori[label_img == 0] = 0
-
-            y, x = util.npmax(map_ori)
-            y = int(float(y) * float(Hr) / float(wsize))
-            x = int(float(x) * float(Wr) / float(wsize))
-            all_peaks.append([x, y])
-        return np.array(all_peaks)
-
-if __name__ == "__main__":
-    hand_estimation = Hand('../model/hand_pose_model.pth')
-
-    # test_image = '../images/hand.jpg'
-    test_image = '../images/hand.jpg'
-    oriImg = cv2.imread(test_image)  # B,G,R order
-    peaks = hand_estimation(oriImg)
-    canvas = util.draw_handpose(oriImg, peaks, True)
-    cv2.imshow('', canvas)
-    cv2.waitKey(0)

idm_masking/openpose/annotator/openpose/model.py

@@ -1,219 +0,0 @@
-import torch
-from collections import OrderedDict
-
-import torch
-import torch.nn as nn
-
-def make_layers(block, no_relu_layers):
-    layers = []
-    for layer_name, v in block.items():
-        if 'pool' in layer_name:
-            layer = nn.MaxPool2d(kernel_size=v[0], stride=v[1],
-                                    padding=v[2])
-            layers.append((layer_name, layer))
-        else:
-            conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1],
-                               kernel_size=v[2], stride=v[3],
-                               padding=v[4])
-            layers.append((layer_name, conv2d))
-            if layer_name not in no_relu_layers:
-                layers.append(('relu_'+layer_name, nn.ReLU(inplace=True)))
-
-    return nn.Sequential(OrderedDict(layers))
-
-class bodypose_model(nn.Module):
-    def __init__(self):
-        super(bodypose_model, self).__init__()
-
-        # these layers have no relu layer
-        no_relu_layers = ['conv5_5_CPM_L1', 'conv5_5_CPM_L2', 'Mconv7_stage2_L1',\
-                          'Mconv7_stage2_L2', 'Mconv7_stage3_L1', 'Mconv7_stage3_L2',\
-                          'Mconv7_stage4_L1', 'Mconv7_stage4_L2', 'Mconv7_stage5_L1',\
-                          'Mconv7_stage5_L2', 'Mconv7_stage6_L1', 'Mconv7_stage6_L1']
-        blocks = {}
-        block0 = OrderedDict([
-                      ('conv1_1', [3, 64, 3, 1, 1]),
-                      ('conv1_2', [64, 64, 3, 1, 1]),
-                      ('pool1_stage1', [2, 2, 0]),
-                      ('conv2_1', [64, 128, 3, 1, 1]),
-                      ('conv2_2', [128, 128, 3, 1, 1]),
-                      ('pool2_stage1', [2, 2, 0]),
-                      ('conv3_1', [128, 256, 3, 1, 1]),
-                      ('conv3_2', [256, 256, 3, 1, 1]),
-                      ('conv3_3', [256, 256, 3, 1, 1]),
-                      ('conv3_4', [256, 256, 3, 1, 1]),
-                      ('pool3_stage1', [2, 2, 0]),
-                      ('conv4_1', [256, 512, 3, 1, 1]),
-                      ('conv4_2', [512, 512, 3, 1, 1]),
-                      ('conv4_3_CPM', [512, 256, 3, 1, 1]),
-                      ('conv4_4_CPM', [256, 128, 3, 1, 1])
-                  ])
-
-
-        # Stage 1
-        block1_1 = OrderedDict([
-                        ('conv5_1_CPM_L1', [128, 128, 3, 1, 1]),
-                        ('conv5_2_CPM_L1', [128, 128, 3, 1, 1]),
-                        ('conv5_3_CPM_L1', [128, 128, 3, 1, 1]),
-                        ('conv5_4_CPM_L1', [128, 512, 1, 1, 0]),
-                        ('conv5_5_CPM_L1', [512, 38, 1, 1, 0])
-                    ])
-
-        block1_2 = OrderedDict([
-                        ('conv5_1_CPM_L2', [128, 128, 3, 1, 1]),
-                        ('conv5_2_CPM_L2', [128, 128, 3, 1, 1]),
-                        ('conv5_3_CPM_L2', [128, 128, 3, 1, 1]),
-                        ('conv5_4_CPM_L2', [128, 512, 1, 1, 0]),
-                        ('conv5_5_CPM_L2', [512, 19, 1, 1, 0])
-                    ])
-        blocks['block1_1'] = block1_1
-        blocks['block1_2'] = block1_2
-
-        self.model0 = make_layers(block0, no_relu_layers)
-
-        # Stages 2 - 6
-        for i in range(2, 7):
-            blocks['block%d_1' % i] = OrderedDict([
-                    ('Mconv1_stage%d_L1' % i, [185, 128, 7, 1, 3]),
-                    ('Mconv2_stage%d_L1' % i, [128, 128, 7, 1, 3]),
-                    ('Mconv3_stage%d_L1' % i, [128, 128, 7, 1, 3]),
-                    ('Mconv4_stage%d_L1' % i, [128, 128, 7, 1, 3]),
-                    ('Mconv5_stage%d_L1' % i, [128, 128, 7, 1, 3]),
-                    ('Mconv6_stage%d_L1' % i, [128, 128, 1, 1, 0]),
-                    ('Mconv7_stage%d_L1' % i, [128, 38, 1, 1, 0])
-                ])
-
-            blocks['block%d_2' % i] = OrderedDict([
-                    ('Mconv1_stage%d_L2' % i, [185, 128, 7, 1, 3]),
-                    ('Mconv2_stage%d_L2' % i, [128, 128, 7, 1, 3]),
-                    ('Mconv3_stage%d_L2' % i, [128, 128, 7, 1, 3]),
-                    ('Mconv4_stage%d_L2' % i, [128, 128, 7, 1, 3]),
-                    ('Mconv5_stage%d_L2' % i, [128, 128, 7, 1, 3]),
-                    ('Mconv6_stage%d_L2' % i, [128, 128, 1, 1, 0]),
-                    ('Mconv7_stage%d_L2' % i, [128, 19, 1, 1, 0])
-                ])
-
-        for k in blocks.keys():
-            blocks[k] = make_layers(blocks[k], no_relu_layers)
-
-        self.model1_1 = blocks['block1_1']
-        self.model2_1 = blocks['block2_1']
-        self.model3_1 = blocks['block3_1']
-        self.model4_1 = blocks['block4_1']
-        self.model5_1 = blocks['block5_1']
-        self.model6_1 = blocks['block6_1']
-
-        self.model1_2 = blocks['block1_2']
-        self.model2_2 = blocks['block2_2']
-        self.model3_2 = blocks['block3_2']
-        self.model4_2 = blocks['block4_2']
-        self.model5_2 = blocks['block5_2']
-        self.model6_2 = blocks['block6_2']
-
-
-    def forward(self, x):
-
-        out1 = self.model0(x)
-
-        out1_1 = self.model1_1(out1)
-        out1_2 = self.model1_2(out1)
-        out2 = torch.cat([out1_1, out1_2, out1], 1)
-
-        out2_1 = self.model2_1(out2)
-        out2_2 = self.model2_2(out2)
-        out3 = torch.cat([out2_1, out2_2, out1], 1)
-
-        out3_1 = self.model3_1(out3)
-        out3_2 = self.model3_2(out3)
-        out4 = torch.cat([out3_1, out3_2, out1], 1)
-
-        out4_1 = self.model4_1(out4)
-        out4_2 = self.model4_2(out4)
-        out5 = torch.cat([out4_1, out4_2, out1], 1)
-
-        out5_1 = self.model5_1(out5)
-        out5_2 = self.model5_2(out5)
-        out6 = torch.cat([out5_1, out5_2, out1], 1)
-
-        out6_1 = self.model6_1(out6)
-        out6_2 = self.model6_2(out6)
-
-        return out6_1, out6_2
-
-class handpose_model(nn.Module):
-    def __init__(self):
-        super(handpose_model, self).__init__()
-
-        # these layers have no relu layer
-        no_relu_layers = ['conv6_2_CPM', 'Mconv7_stage2', 'Mconv7_stage3',\
-                          'Mconv7_stage4', 'Mconv7_stage5', 'Mconv7_stage6']
-        # stage 1
-        block1_0 = OrderedDict([
-                ('conv1_1', [3, 64, 3, 1, 1]),
-                ('conv1_2', [64, 64, 3, 1, 1]),
-                ('pool1_stage1', [2, 2, 0]),
-                ('conv2_1', [64, 128, 3, 1, 1]),
-                ('conv2_2', [128, 128, 3, 1, 1]),
-                ('pool2_stage1', [2, 2, 0]),
-                ('conv3_1', [128, 256, 3, 1, 1]),
-                ('conv3_2', [256, 256, 3, 1, 1]),
-                ('conv3_3', [256, 256, 3, 1, 1]),
-                ('conv3_4', [256, 256, 3, 1, 1]),
-                ('pool3_stage1', [2, 2, 0]),
-                ('conv4_1', [256, 512, 3, 1, 1]),
-                ('conv4_2', [512, 512, 3, 1, 1]),
-                ('conv4_3', [512, 512, 3, 1, 1]),
-                ('conv4_4', [512, 512, 3, 1, 1]),
-                ('conv5_1', [512, 512, 3, 1, 1]),
-                ('conv5_2', [512, 512, 3, 1, 1]),
-                ('conv5_3_CPM', [512, 128, 3, 1, 1])
-            ])
-
-        block1_1 = OrderedDict([
-            ('conv6_1_CPM', [128, 512, 1, 1, 0]),
-            ('conv6_2_CPM', [512, 22, 1, 1, 0])
-        ])
-
-        blocks = {}
-        blocks['block1_0'] = block1_0
-        blocks['block1_1'] = block1_1
-
-        # stage 2-6
-        for i in range(2, 7):
-            blocks['block%d' % i] = OrderedDict([
-                    ('Mconv1_stage%d' % i, [150, 128, 7, 1, 3]),
-                    ('Mconv2_stage%d' % i, [128, 128, 7, 1, 3]),
-                    ('Mconv3_stage%d' % i, [128, 128, 7, 1, 3]),
-                    ('Mconv4_stage%d' % i, [128, 128, 7, 1, 3]),
-                    ('Mconv5_stage%d' % i, [128, 128, 7, 1, 3]),
-                    ('Mconv6_stage%d' % i, [128, 128, 1, 1, 0]),
-                    ('Mconv7_stage%d' % i, [128, 22, 1, 1, 0])
-                ])
-
-        for k in blocks.keys():
-            blocks[k] = make_layers(blocks[k], no_relu_layers)
-
-        self.model1_0 = blocks['block1_0']
-        self.model1_1 = blocks['block1_1']
-        self.model2 = blocks['block2']
-        self.model3 = blocks['block3']
-        self.model4 = blocks['block4']
-        self.model5 = blocks['block5']
-        self.model6 = blocks['block6']
-
-    def forward(self, x):
-        out1_0 = self.model1_0(x)
-        out1_1 = self.model1_1(out1_0)
-        concat_stage2 = torch.cat([out1_1, out1_0], 1)
-        out_stage2 = self.model2(concat_stage2)
-        concat_stage3 = torch.cat([out_stage2, out1_0], 1)
-        out_stage3 = self.model3(concat_stage3)
-        concat_stage4 = torch.cat([out_stage3, out1_0], 1)
-        out_stage4 = self.model4(concat_stage4)
-        concat_stage5 = torch.cat([out_stage4, out1_0], 1)
-        out_stage5 = self.model5(concat_stage5)
-        concat_stage6 = torch.cat([out_stage5, out1_0], 1)
-        out_stage6 = self.model6(concat_stage6)
-        return out_stage6
-
-

idm_masking/openpose/annotator/openpose/util.py

@@ -1,297 +0,0 @@
-import math
-import numpy as np
-import matplotlib
-import cv2
-
-
-eps = 0.01
-
-
-def smart_resize(x, s):
-    Ht, Wt = s
-    if x.ndim == 2:
-        Ho, Wo = x.shape
-        Co = 1
-    else:
-        Ho, Wo, Co = x.shape
-    if Co == 3 or Co == 1:
-        k = float(Ht + Wt) / float(Ho + Wo)
-        return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4)
-    else:
-        return np.stack([smart_resize(x[:, :, i], s) for i in range(Co)], axis=2)
-
-
-def smart_resize_k(x, fx, fy):
-    if x.ndim == 2:
-        Ho, Wo = x.shape
-        Co = 1
-    else:
-        Ho, Wo, Co = x.shape
-    Ht, Wt = Ho * fy, Wo * fx
-    if Co == 3 or Co == 1:
-        k = float(Ht + Wt) / float(Ho + Wo)
-        return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4)
-    else:
-        return np.stack([smart_resize_k(x[:, :, i], fx, fy) for i in range(Co)], axis=2)
-
-
-def padRightDownCorner(img, stride, padValue):
-    h = img.shape[0]
-    w = img.shape[1]
-
-    pad = 4 * [None]
-    pad[0] = 0 # up
-    pad[1] = 0 # left
-    pad[2] = 0 if (h % stride == 0) else stride - (h % stride) # down
-    pad[3] = 0 if (w % stride == 0) else stride - (w % stride) # right
-
-    img_padded = img
-    pad_up = np.tile(img_padded[0:1, :, :]*0 + padValue, (pad[0], 1, 1))
-    img_padded = np.concatenate((pad_up, img_padded), axis=0)
-    pad_left = np.tile(img_padded[:, 0:1, :]*0 + padValue, (1, pad[1], 1))
-    img_padded = np.concatenate((pad_left, img_padded), axis=1)
-    pad_down = np.tile(img_padded[-2:-1, :, :]*0 + padValue, (pad[2], 1, 1))
-    img_padded = np.concatenate((img_padded, pad_down), axis=0)
-    pad_right = np.tile(img_padded[:, -2:-1, :]*0 + padValue, (1, pad[3], 1))
-    img_padded = np.concatenate((img_padded, pad_right), axis=1)
-
-    return img_padded, pad
-
-
-def transfer(model, model_weights):
-    transfered_model_weights = {}
-    for weights_name in model.state_dict().keys():
-        transfered_model_weights[weights_name] = model_weights['.'.join(weights_name.split('.')[1:])]
-    return transfered_model_weights
-
-
-def draw_bodypose(canvas, candidate, subset):
-    H, W, C = canvas.shape
-    candidate = np.array(candidate)
-    subset = np.array(subset)
-
-    stickwidth = 4
-
-    limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
-               [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
-               [1, 16], [16, 18], [3, 17], [6, 18]]
-
-    colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
-              [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
-              [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
-
-    for i in range(17):
-        for n in range(len(subset)):
-            index = subset[n][np.array(limbSeq[i]) - 1]
-            if -1 in index:
-                continue
-            Y = candidate[index.astype(int), 0] * float(W)
-            X = candidate[index.astype(int), 1] * float(H)
-            mX = np.mean(X)
-            mY = np.mean(Y)
-            length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
-            angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
-            polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
-            cv2.fillConvexPoly(canvas, polygon, colors[i])
-
-    canvas = (canvas * 0.6).astype(np.uint8)
-
-    for i in range(18):
-        for n in range(len(subset)):
-            index = int(subset[n][i])
-            if index == -1:
-                continue
-            x, y = candidate[index][0:2]
-            x = int(x * W)
-            y = int(y * H)
-            cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1)
-
-    return canvas
-
-
-def draw_handpose(canvas, all_hand_peaks):
-    H, W, C = canvas.shape
-
-    edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \
-             [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]]
-
-    for peaks in all_hand_peaks:
-        peaks = np.array(peaks)
-
-        for ie, e in enumerate(edges):
-            x1, y1 = peaks[e[0]]
-            x2, y2 = peaks[e[1]]
-            x1 = int(x1 * W)
-            y1 = int(y1 * H)
-            x2 = int(x2 * W)
-            y2 = int(y2 * H)
-            if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
-                cv2.line(canvas, (x1, y1), (x2, y2), matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * 255, thickness=2)
-
-        for i, keyponit in enumerate(peaks):
-            x, y = keyponit
-            x = int(x * W)
-            y = int(y * H)
-            if x > eps and y > eps:
-                cv2.circle(canvas, (x, y), 4, (0, 0, 255), thickness=-1)
-    return canvas
-
-
-def draw_facepose(canvas, all_lmks):
-    H, W, C = canvas.shape
-    for lmks in all_lmks:
-        lmks = np.array(lmks)
-        for lmk in lmks:
-            x, y = lmk
-            x = int(x * W)
-            y = int(y * H)
-            if x > eps and y > eps:
-                cv2.circle(canvas, (x, y), 3, (255, 255, 255), thickness=-1)
-    return canvas
-
-
-# detect hand according to body pose keypoints
-# please refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/hand/handDetector.cpp
-def handDetect(candidate, subset, oriImg):
-    # right hand: wrist 4, elbow 3, shoulder 2
-    # left hand: wrist 7, elbow 6, shoulder 5
-    ratioWristElbow = 0.33
-    detect_result = []
-    image_height, image_width = oriImg.shape[0:2]
-    for person in subset.astype(int):
-        # if any of three not detected
-        has_left = np.sum(person[[5, 6, 7]] == -1) == 0
-        has_right = np.sum(person[[2, 3, 4]] == -1) == 0
-        if not (has_left or has_right):
-            continue
-        hands = []
-        #left hand
-        if has_left:
-            left_shoulder_index, left_elbow_index, left_wrist_index = person[[5, 6, 7]]
-            x1, y1 = candidate[left_shoulder_index][:2]
-            x2, y2 = candidate[left_elbow_index][:2]
-            x3, y3 = candidate[left_wrist_index][:2]
-            hands.append([x1, y1, x2, y2, x3, y3, True])
-        # right hand
-        if has_right:
-            right_shoulder_index, right_elbow_index, right_wrist_index = person[[2, 3, 4]]
-            x1, y1 = candidate[right_shoulder_index][:2]
-            x2, y2 = candidate[right_elbow_index][:2]
-            x3, y3 = candidate[right_wrist_index][:2]
-            hands.append([x1, y1, x2, y2, x3, y3, False])
-
-        for x1, y1, x2, y2, x3, y3, is_left in hands:
-            # pos_hand = pos_wrist + ratio * (pos_wrist - pos_elbox) = (1 + ratio) * pos_wrist - ratio * pos_elbox
-            # handRectangle.x = posePtr[wrist*3] + ratioWristElbow * (posePtr[wrist*3] - posePtr[elbow*3]);
-            # handRectangle.y = posePtr[wrist*3+1] + ratioWristElbow * (posePtr[wrist*3+1] - posePtr[elbow*3+1]);
-            # const auto distanceWristElbow = getDistance(poseKeypoints, person, wrist, elbow);
-            # const auto distanceElbowShoulder = getDistance(poseKeypoints, person, elbow, shoulder);
-            # handRectangle.width = 1.5f * fastMax(distanceWristElbow, 0.9f * distanceElbowShoulder);
-            x = x3 + ratioWristElbow * (x3 - x2)
-            y = y3 + ratioWristElbow * (y3 - y2)
-            distanceWristElbow = math.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2)
-            distanceElbowShoulder = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
-            width = 1.5 * max(distanceWristElbow, 0.9 * distanceElbowShoulder)
-            # x-y refers to the center --> offset to topLeft point
-            # handRectangle.x -= handRectangle.width / 2.f;
-            # handRectangle.y -= handRectangle.height / 2.f;
-            x -= width / 2
-            y -= width / 2  # width = height
-            # overflow the image
-            if x < 0: x = 0
-            if y < 0: y = 0
-            width1 = width
-            width2 = width
-            if x + width > image_width: width1 = image_width - x
-            if y + width > image_height: width2 = image_height - y
-            width = min(width1, width2)
-            # the max hand box value is 20 pixels
-            if width >= 20:
-                detect_result.append([int(x), int(y), int(width), is_left])
-
-    '''
-    return value: [[x, y, w, True if left hand else False]].
-    width=height since the network require squared input.
-    x, y is the coordinate of top left 
-    '''
-    return detect_result
-
-
-# Written by Lvmin
-def faceDetect(candidate, subset, oriImg):
-    # left right eye ear 14 15 16 17
-    detect_result = []
-    image_height, image_width = oriImg.shape[0:2]
-    for person in subset.astype(int):
-        has_head = person[0] > -1
-        if not has_head:
-            continue
-
-        has_left_eye = person[14] > -1
-        has_right_eye = person[15] > -1
-        has_left_ear = person[16] > -1
-        has_right_ear = person[17] > -1
-
-        if not (has_left_eye or has_right_eye or has_left_ear or has_right_ear):
-            continue
-
-        head, left_eye, right_eye, left_ear, right_ear = person[[0, 14, 15, 16, 17]]
-
-        width = 0.0
-        x0, y0 = candidate[head][:2]
-
-        if has_left_eye:
-            x1, y1 = candidate[left_eye][:2]
-            d = max(abs(x0 - x1), abs(y0 - y1))
-            width = max(width, d * 3.0)
-
-        if has_right_eye:
-            x1, y1 = candidate[right_eye][:2]
-            d = max(abs(x0 - x1), abs(y0 - y1))
-            width = max(width, d * 3.0)
-
-        if has_left_ear:
-            x1, y1 = candidate[left_ear][:2]
-            d = max(abs(x0 - x1), abs(y0 - y1))
-            width = max(width, d * 1.5)
-
-        if has_right_ear:
-            x1, y1 = candidate[right_ear][:2]
-            d = max(abs(x0 - x1), abs(y0 - y1))
-            width = max(width, d * 1.5)
-
-        x, y = x0, y0
-
-        x -= width
-        y -= width
-
-        if x < 0:
-            x = 0
-
-        if y < 0:
-            y = 0
-
-        width1 = width * 2
-        width2 = width * 2
-
-        if x + width > image_width:
-            width1 = image_width - x
-
-        if y + width > image_height:
-            width2 = image_height - y
-
-        width = min(width1, width2)
-
-        if width >= 20:
-            detect_result.append([int(x), int(y), int(width)])
-
-    return detect_result
-
-
-# get max index of 2d array
-def npmax(array):
-    arrayindex = array.argmax(1)
-    arrayvalue = array.max(1)
-    i = arrayvalue.argmax()
-    j = arrayindex[i]
-    return i, j

idm_masking/openpose/annotator/util.py

@@ -1,100 +0,0 @@
-import random
-
-import numpy as np
-import cv2
-import os
-from pathlib import Path
-
-PROJECT_ROOT = Path(__file__).absolute().parents[3].absolute()
-
-annotator_ckpts_path = os.path.join(PROJECT_ROOT, 'ckpt/openpose/ckpts')
-# print(annotator_ckpts_path)
-
-def HWC3(x):
-    assert x.dtype == np.uint8
-    if x.ndim == 2:
-        x = x[:, :, None]
-    assert x.ndim == 3
-    H, W, C = x.shape
-    assert C == 1 or C == 3 or C == 4
-    if C == 3:
-        return x
-    if C == 1:
-        return np.concatenate([x, x, x], axis=2)
-    if C == 4:
-        color = x[:, :, 0:3].astype(np.float32)
-        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
-        y = color * alpha + 255.0 * (1.0 - alpha)
-        y = y.clip(0, 255).astype(np.uint8)
-        return y
-
-
-def resize_image(input_image, resolution):
-    H, W, C = input_image.shape
-    H = float(H)
-    W = float(W)
-    k = float(resolution) / min(H, W)
-    H *= k
-    W *= k
-    H = int(np.round(H / 64.0)) * 64
-    W = int(np.round(W / 64.0)) * 64
-    img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
-    return img
-
-
-def nms(x, t, s):
-    x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
-
-    f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
-    f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
-    f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
-    f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
-
-    y = np.zeros_like(x)
-
-    for f in [f1, f2, f3, f4]:
-        np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
-
-    z = np.zeros_like(y, dtype=np.uint8)
-    z[y > t] = 255
-    return z
-
-
-def make_noise_disk(H, W, C, F):
-    noise = np.random.uniform(low=0, high=1, size=((H // F) + 2, (W // F) + 2, C))
-    noise = cv2.resize(noise, (W + 2 * F, H + 2 * F), interpolation=cv2.INTER_CUBIC)
-    noise = noise[F: F + H, F: F + W]
-    noise -= np.min(noise)
-    noise /= np.max(noise)
-    if C == 1:
-        noise = noise[:, :, None]
-    return noise
-
-
-def min_max_norm(x):
-    x -= np.min(x)
-    x /= np.maximum(np.max(x), 1e-5)
-    return x
-
-
-def safe_step(x, step=2):
-    y = x.astype(np.float32) * float(step + 1)
-    y = y.astype(np.int32).astype(np.float32) / float(step)
-    return y
-
-
-def img2mask(img, H, W, low=10, high=90):
-    assert img.ndim == 3 or img.ndim == 2
-    assert img.dtype == np.uint8
-
-    if img.ndim == 3:
-        y = img[:, :, random.randrange(0, img.shape[2])]
-    else:
-        y = img
-
-    y = cv2.resize(y, (W, H), interpolation=cv2.INTER_CUBIC)
-
-    if random.uniform(0, 1) < 0.5:
-        y = 255 - y
-
-    return y < np.percentile(y, random.randrange(low, high))

idm_masking/openpose/run_openpose.py

@@ -1,84 +0,0 @@
-import pdb
-
-# import config
-from pathlib import Path
-import sys
-
-PROJECT_ROOT = Path(__file__).absolute().parents[0].absolute()
-sys.path.insert(0, str(PROJECT_ROOT))
-import os
-
-import cv2
-import einops
-import numpy as np
-import random
-import time
-import json
-
-# from pytorch_lightning import seed_everything
-from annotator.util import resize_image, HWC3
-from annotator.openpose import OpenposeDetector
-
-import argparse
-from PIL import Image
-import torch
-import pdb
-
-# os.environ['CUDA_VISIBLE_DEVICES'] = '0,1,2,3'
-
-class OpenPose:
-    def __init__(self, gpu_id: int):
-        self.gpu_id = gpu_id
-        torch.cuda.set_device(gpu_id)
-        self.preprocessor = OpenposeDetector()
-
-    def __call__(self, input_image, resolution=384):
-        torch.cuda.set_device(self.gpu_id)
-        if isinstance(input_image, Image.Image):
-            input_image = np.asarray(input_image)
-        elif type(input_image) == str:
-            input_image = np.asarray(Image.open(input_image))
-        else:
-            raise ValueError
-        with torch.no_grad():
-            input_image = HWC3(input_image)
-            input_image = resize_image(input_image, resolution)
-            H, W, C = input_image.shape
-            assert (H == 512 and W == 384), 'Incorrect input image shape'
-            pose, detected_map = self.preprocessor(input_image, hand_and_face=False)
-
-            candidate = pose['bodies']['candidate']
-            subset = pose['bodies']['subset'][0][:18]
-            for i in range(18):
-                if subset[i] == -1:
-                    candidate.insert(i, [0, 0])
-                    for j in range(i, 18):
-                        if(subset[j]) != -1:
-                            subset[j] += 1
-                elif subset[i] != i:
-                    candidate.pop(i)
-                    for j in range(i, 18):
-                        if(subset[j]) != -1:
-                            subset[j] -= 1
-
-            candidate = candidate[:18]
-
-            for i in range(18):
-                candidate[i][0] *= 384
-                candidate[i][1] *= 512
-
-            keypoints = {"pose_keypoints_2d": candidate}
-            # with open("/home/aigc/ProjectVTON/OpenPose/keypoints/keypoints.json", "w") as f:
-            #     json.dump(keypoints, f)
-            #
-            # # print(candidate)
-            # output_image = cv2.resize(cv2.cvtColor(detected_map, cv2.COLOR_BGR2RGB), (768, 1024))
-            # cv2.imwrite('/home/aigc/ProjectVTON/OpenPose/keypoints/out_pose.jpg', output_image)
-
-        return keypoints
-
-
-if __name__ == '__main__':
-
-    model = OpenPose()
-    model('./images/bad_model.jpg')

image_stack.py

@@ -0,0 +1,203 @@
+#!/usr/bin/python3
+
+from PIL import Image, ImageOps, ImageSequence, ImageFile
+from PIL.PngImagePlugin import PngInfo
+import cv2
+import hashlib
+import json
+import logging
+import math
+import numpy as np
+import os
+import random
+import safetensors.torch
+import sys
+import time
+import torch
+import traceback
+
+
+def load_image(path):
+
+    return torch.from_numpy(cv2.imread(
+        path, cv2.IMREAD_COLOR)).to(dtype=torch.float32) / 255.0
+
+
+def do_stack(img1, img2):
+
+    dim = max(max(img1.shape[0], img2.shape[0]), img1.shape[1] + img2.shape[1])
+
+    out = torch.zeros((dim, dim, 3), dtype=img1.dtype, device=img1.device) + 1
+
+    diff1 = (out.shape[0] - img1.shape[0]) // 2
+    diff2 = (out.shape[0] - img2.shape[0]) // 2
+
+    part0 = 0
+    part1 = img1.shape[1]
+    part2 = img2.shape[1] + img1.shape[1]
+
+    out[diff1:diff1 + img1.shape[0], part0:part1, :] = img1
+    out[diff2:diff2 + img2.shape[0], part1:part2, :] = img2
+
+    return out
+
+
+def save_image(image, outpath):
+
+    cv2.imwrite(outpath,
+                (image * 255).to(dtype=torch.uint8).detach().cpu().numpy())
+
+
+class H_Stack_Images:
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image_L": ("IMAGE", ),
+                "image_R": ("IMAGE", ),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE", )
+    FUNCTION = "test"
+    CATEGORY = "TRI3D"
+
+    def test(self, image_L, image_R):
+
+        return (do_stack(img1=image_L[0], img2=image_R[0]).unsqueeze(0), )
+
+
+class SaveImage_absolute:
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE", {
+                    "tooltip": "The images to save."
+                }),
+                "absolute_filename": ("STRING", {
+                    "default":
+                    "image.png",
+                    "tooltip":
+                    "The absolute path to the file to save."
+                })
+            },
+        }
+
+    RETURN_TYPES = ("STRING", )
+    RETURN_NAMES = ("text to control order", )
+    FUNCTION = "save_images"
+
+    OUTPUT_NODE = True
+
+    CATEGORY = "image"
+    DESCRIPTION = "Saves the input images to an absolute path."
+
+    def save_images(self, images, absolute_filename):
+        i = 255.0 * images[0].cpu().numpy()
+        img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+        img.save(absolute_filename)
+
+        return (absolute_filename, )
+
+
+class SaveText_absolute:
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "text": ("STRING", {
+                    "multiline": True,
+                    "dynamicPrompts": True,
+                    "tooltip": "Text to be saved to the file."
+                }),
+                "absolute_filename": ("STRING", {
+                    "default":
+                    "image.txt",
+                    "tooltip":
+                    "The absolute path to the file to save."
+                })
+            },
+            "optional": {
+                "text_opt": ("STRING", {
+                    "multiline":
+                    True,
+                    "dynamicPrompts":
+                    True,
+                    "tooltip":
+                    "Text to provide order when necessary (to create work files after txt files)."
+                }),
+            }
+        }
+
+    RETURN_TYPES = ("STRING", )
+    RETURN_NAMES = ("same text as input", )
+    FUNCTION = "save_text"
+
+    OUTPUT_NODE = True
+
+    CATEGORY = "text"
+    DESCRIPTION = "Saves the input text to an absolute path."
+
+    def save_text(self, text, absolute_filename, text_opt=''):
+        open(absolute_filename, "w").write(text)
+        return (text, )
+
+
+class Wait_And_Read_File:
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "absolute_filename": ("STRING", {
+                    "default":
+                    "image.txt",
+                    "tooltip":
+                    "The absolute path to the file to read."
+                })
+            },
+            "optional": {
+                "text": ("STRING", {
+                    "multiline":
+                    True,
+                    "dynamicPrompts":
+                    True,
+                    "tooltip":
+                    "Text to provide order when necessary (to wait on done file)."
+                }),
+            }
+        }
+
+    RETURN_TYPES = ("STRING", )
+    RETURN_NAMES = ("text from file", )
+    FUNCTION = "read_text"
+
+    OUTPUT_NODE = True
+
+    CATEGORY = "text"
+    DESCRIPTION = "Saves the input text to an absolute path."
+
+    def read_text(self, absolute_filename, text=''):
+        while not os.path.exists(absolute_filename):
+            time.sleep(0.1)
+
+        res = open(absolute_filename, "r").read()
+        os.unlink(absolute_filename)
+
+        return (res, )

mask_area_percentage.py

@@ -0,0 +1,68 @@
+import os
+import cv2
+import numpy as np
+import torch
+
+class TRI3D_MaskAreaPercentage:
+    """
+    ComfyUI node that calculates the percentage of white pixels in an image
+    relative to the total image area.
+    """
+    
+    def from_torch_image(self, image):
+        """Convert a torch tensor image to numpy array for OpenCV processing"""
+        image = image.cpu().numpy() * 255.0
+        image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "threshold": ("INT", {"default": 127, "min": 0, "max": 255, "step": 1}),
+            },
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("FLOAT", "INT", "INT",)
+    RETURN_NAMES = ("percentage", "white_pixels", "total_pixels",)
+    CATEGORY = "TRI3D"
+
+    def run(self, image, threshold=127):
+        # Convert Torch image to OpenCV format
+        cv_image = self.from_torch_image(image)
+        
+        # Remove batch dimension if present
+        if len(cv_image.shape) == 4:
+            cv_image = cv_image[0]
+        
+        # Convert to grayscale if it's a color image
+        if len(cv_image.shape) == 3 and cv_image.shape[2] > 1:
+            gray_image = cv2.cvtColor(cv_image, cv2.COLOR_RGB2GRAY)
+        else:
+            gray_image = cv_image[:, :, 0]
+        
+        # Calculate total number of pixels
+        total_pixels = gray_image.shape[0] * gray_image.shape[1]
+        
+        # Count white pixels (pixels with values above threshold)
+        _, binary_image = cv2.threshold(gray_image, threshold, 255, cv2.THRESH_BINARY)
+        white_pixels = cv2.countNonZero(binary_image)
+        
+        # Calculate percentage of white pixels
+        percentage = (white_pixels / total_pixels) * 100.0
+        
+        return (percentage, white_pixels, total_pixels,)
+
+# # Node registration for ComfyUI
+# NODE_CLASS_MAPPINGS = {
+#     "TRI3D_MaskAreaPercentage": TRI3D_MaskAreaPercentage
+# }
+
+# NODE_DISPLAY_NAME_MAPPINGS = {
+#     "TRI3D_MaskAreaPercentage": "TRI3D Mask Area Percentage"
+# } 

nsfw.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+from weakref import ref as WeakRef
+from pathlib import Path
+from tqdm import tqdm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from transformers import CLIPImageProcessor, CLIPConfig, CLIPVisionModel, PreTrainedModel
+from kornia.filters import box_blur
+
+
+
+def cosine_similarity(image_embeds: Tensor, text_embeds: Tensor):
+    if image_embeds.dim() == 2 and text_embeds.dim() == 2:
+        image_embeds = image_embeds.unsqueeze(1)
+    return F.cosine_similarity(image_embeds, text_embeds, dim=-1)
+
+
+class CLIPSafetyChecker(PreTrainedModel):
+    # https://huggingface.co/CompVis/stable-diffusion-safety-checker
+    # Adapted from:
+    # https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/safety_checker.py
+
+    config_class = CLIPConfig
+    _no_split_modules = ["CLIPEncoderLayer"]
+
+    def __init__(self, config: CLIPConfig):
+        super().__init__(config)
+        projdim = config.projection_dim
+
+        self.vision_model = CLIPVisionModel(config.vision_config)
+        self.visual_projection = nn.Linear(config.vision_config.hidden_size, projdim, bias=False)
+
+        self.concept_embeds = nn.Parameter(torch.ones(17, projdim), requires_grad=False)
+        self.special_care_embeds = nn.Parameter(torch.ones(3, projdim), requires_grad=False)
+        self.concept_embeds_weights = nn.Parameter(torch.ones(17), requires_grad=False)
+        self.special_care_embeds_weights = nn.Parameter(torch.ones(3), requires_grad=False)
+
+    def forward(self, clip_input, images: Tensor, sensitivity: float, alternate_image: Tensor):
+        with torch.no_grad():
+            image_batch = self.vision_model(clip_input)[1]
+            image_embeds = self.visual_projection(image_batch)
+            sensitivity = -0.1 + 0.14 * sensitivity
+
+            special_cos_dist = cosine_similarity(image_embeds, self.special_care_embeds)
+            special_scores_threshold = self.special_care_embeds_weights.unsqueeze(0)
+            special_scores = special_cos_dist - special_scores_threshold + sensitivity
+
+            if torch.any(special_scores > 0):
+                sensitivity = sensitivity + 0.01
+
+            cos_dist = cosine_similarity(image_embeds, self.concept_embeds)
+            concept_threshold = self.concept_embeds_weights.unsqueeze(0)
+            concept_scores = cos_dist - concept_threshold + sensitivity
+
+            is_nsfw = [torch.any(concept_scores[i] > 0) for i in range(concept_scores.shape[0])]
+            is_nsfw = [x.item() for x in is_nsfw]
+            return self.filter_images(images, alternate_image, is_nsfw)
+
+    def filter_images(self, images: Tensor, alternate_image: Tensor, is_nsfw: list[bool]):
+        if not any(is_nsfw):
+            return images
+
+        images = images.clone()
+        for idx, nsfw in enumerate(is_nsfw):
+            if nsfw:
+                # Resize alternate image to match original image dimensions
+                resized_alternate = F.interpolate(
+                    alternate_image[idx:idx+1],  # Add batch dimension
+                    size=(images[idx].shape[1], images[idx].shape[2]),  # Target height, width
+                    mode='bilinear',
+                    align_corners=False
+                )
+                images[idx] = resized_alternate.squeeze(0)  # Remove batch dimension
+        return images
+
+
+class CachedModels:
+    _instance: WeakRef | None = None
+
+    def __init__(self):
+        model_dir = Path(__file__).parent / "safetychecker"
+        model_file = model_dir / "model.safetensors"
+        if not model_file.exists():
+            self.download(
+                "https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/refs%2Fpr%2F41/model.safetensors",
+                target=model_file,
+            )
+        self.feature_extractor = CLIPImageProcessor.from_pretrained(model_dir)
+        self.safety_checker = CLIPSafetyChecker.from_pretrained(model_dir)
+
+    @classmethod
+    def load(cls):
+        models = cls._instance and cls._instance()
+        if models is None:
+            models = cls()
+            cls._instance = WeakRef(models)
+        return models
+
+    def download(self, url: str, target: Path):
+        import requests
+
+        try:
+            target_temp = target.with_suffix(".download")
+            with requests.get(url, stream=True) as response:
+                text = "NSFWFilter model download"
+                total = int(response.headers.get("content-length", 0))
+                pbar = tqdm(None, total=total, unit="b", unit_scale=True, desc=text)
+                with open(target_temp, "wb") as f:
+                    for chunk in response.iter_content(chunk_size=8192):
+                        f.write(chunk)
+                        pbar.update(len(chunk))
+                pbar.close()
+            target_temp.rename(target)
+        except Exception as e:
+            raise RuntimeError(
+                f"NSFWFilter: Failed to download safety-checker model from {url} to target location {target}: {e}"
+            ) from e
+
+
+
+def to_bchw(image: torch.Tensor):
+    if image.ndim == 3:
+        image = image.unsqueeze(0)
+    return image.movedim(-1, 1)
+
+
+def to_bhwc(image: torch.Tensor):
+    return image.movedim(1, -1)
+
+
+def mask_batch(mask: torch.Tensor):
+    if mask.ndim == 2:
+        mask = mask.unsqueeze(0)
+    return mask
+
+
+class TRI3DNSFWFilter:
+    models: CachedModels
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "alternate_image": ("IMAGE",),
+                "sensitivity": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.10}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "check"
+    CATEGORY = "TRI3D NSFW"
+
+    def __init__(self):
+        self.models = CachedModels.load()
+
+    def check(self, image, alternate_image,sensitivity):
+        image = to_bchw(image)
+        alternate_image = to_bchw(alternate_image)
+        input = self.models.feature_extractor(image, do_rescale=False, return_tensors="pt")
+        filtered = self.models.safety_checker(
+            images=image, clip_input=input.pixel_values, sensitivity=sensitivity, alternate_image=alternate_image
+        )
+        return (to_bhwc(filtered),)

remove_small_mask_islands.py

@@ -0,0 +1,117 @@
+import os
+import cv2
+import numpy as np
+import torch
+
+class TRI3D_RemoveSmallMaskIslands:
+    """
+    ComfyUI node that removes small islands of white pixels from a mask image
+    based on a specified area threshold.
+    """
+    
+    def from_torch_image(self, image):
+        """Convert a torch tensor image to numpy array for OpenCV processing"""
+        image = image.cpu().numpy() * 255.0
+        image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    def to_torch_image(self, image):
+        """Convert numpy array back to torch tensor format"""
+        image = image.astype(dtype=np.float32)
+        image /= 255.0
+        image = torch.from_numpy(image)
+        return image
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE", ),
+                "min_island_area": ("INT", {"default": 100, "min": 1, "max": 10000, "step": 10}),
+                "invert": ("BOOLEAN", {"default": False}),
+            },
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE", )
+    CATEGORY = "TRI3D"
+
+    def run(self, image, min_island_area, invert):
+        # Convert Torch image to OpenCV format
+        cv_image = self.from_torch_image(image)
+        
+        # Remove batch dimension if present
+        if len(cv_image.shape) == 4:
+            cv_image = cv_image[0]
+        
+        # Make a copy to work with
+        result_image = cv_image.copy()
+        
+        # Process each channel (if grayscale, it will just be one iteration)
+        height, width = cv_image.shape[:2]
+        
+        # If the image has 3 channels (RGB), convert to grayscale for contour detection
+        if len(cv_image.shape) == 3 and cv_image.shape[2] == 3:
+            # Convert to grayscale for processing
+            gray = cv2.cvtColor(cv_image, cv2.COLOR_RGB2GRAY)
+        else:
+            # Use the first channel if it's already grayscale or has alpha
+            gray = cv_image[:, :, 0]
+        
+        # Invert if needed (to work with black islands instead of white)
+        if invert:
+            gray = 255 - gray
+            
+        # Create binary image
+        _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
+        
+        # Find contours in the binary image
+        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+        # Create a blank mask for the cleaned image
+        clean_mask = np.zeros((height, width), dtype=np.uint8)
+        
+        # Draw only contours with area greater than the threshold
+        for contour in contours:
+            area = cv2.contourArea(contour)
+            if area >= min_island_area:
+                cv2.drawContours(clean_mask, [contour], 0, 255, -1)
+        
+        # Invert back if needed
+        if invert:
+            clean_mask = 255 - clean_mask
+        
+        # Apply the clean mask to each channel of the original image
+        if len(cv_image.shape) == 3 and cv_image.shape[2] == 3:
+            # RGB image
+            for i in range(3):
+                result_image[:, :, i] = cv2.bitwise_and(cv_image[:, :, i], clean_mask)
+        elif len(cv_image.shape) == 3 and cv_image.shape[2] == 4:
+            # RGBA image
+            for i in range(4):
+                result_image[:, :, i] = cv2.bitwise_and(cv_image[:, :, i], clean_mask)
+        else:
+            # Single channel image
+            result_image = cv2.bitwise_and(cv_image, clean_mask)
+            # Reshape to match expected dimensions
+            result_image = result_image.reshape(height, width, 1)
+        
+        # Convert back to torch format
+        torch_image = self.to_torch_image(result_image)
+        
+        # Add batch dimension back
+        torch_image = torch_image.unsqueeze(0)
+        
+        return (torch_image,)
+
+# # Node registration for ComfyUI
+# NODE_CLASS_MAPPINGS = {
+#     "TRI3D_RemoveSmallMaskIslands": TRI3D_RemoveSmallMaskIslands
+# }
+
+# NODE_DISPLAY_NAME_MAPPINGS = {
+#     "TRI3D_RemoveSmallMaskIslands": "TRI3D Remove Small Mask Islands"
+# }

safetychecker/config.json

@@ -0,0 +1,171 @@
+{
+  "_name_or_path": "clip-vit-large-patch14/",
+  "architectures": [
+    "SafetyChecker"
+  ],
+  "initializer_factor": 1.0,
+  "logit_scale_init_value": 2.6592,
+  "model_type": "clip",
+  "projection_dim": 768,
+  "text_config": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": null,
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "bos_token_id": 0,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dropout": 0.0,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 2,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "quick_gelu",
+    "hidden_size": 768,
+	"id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_factor": 1.0,
+    "initializer_range": 0.02,
+    "intermediate_size": 3072,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-05,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 77,
+    "min_length": 0,
+    "model_type": "clip_text_model",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 12,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 12,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": 1,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "transformers_version": "4.21.0.dev0",
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "vocab_size": 49408
+  },
+  "text_config_dict": {
+    "hidden_size": 768,
+    "intermediate_size": 3072,
+    "num_attention_heads": 12,
+    "num_hidden_layers": 12
+  },
+  "torch_dtype": "float32",
+  "transformers_version": null,
+  "vision_config": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": null,
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dropout": 0.0,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "quick_gelu",
+    "hidden_size": 1024,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_size": 224,
+    "initializer_factor": 1.0,
+    "initializer_range": 0.02,
+    "intermediate_size": 4096,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-05,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "clip_vision_model",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 16,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 24,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 14,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": null,
+    "torchscript": false,
+    "transformers_version": "4.21.0.dev0",
+    "typical_p": 1.0,
+    "use_bfloat16": false
+  },
+  "vision_config_dict": {
+    "hidden_size": 1024,
+    "intermediate_size": 4096,
+    "num_attention_heads": 16,
+    "num_hidden_layers": 24,
+    "patch_size": 14
+  }
+}

safetychecker/preprocessor_config.json

@@ -0,0 +1,20 @@
+{
+  "crop_size": 224,
+  "do_center_crop": true,
+  "do_convert_rgb": true,
+  "do_normalize": true,
+  "do_resize": true,
+  "feature_extractor_type": "CLIPFeatureExtractor",
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "resample": 3,
+  "size": 224
+}

smart_box.py

@@ -0,0 +1,927 @@
+import numpy as np
+import torch
+import json
+
+import cv2
+
+#  {0,  "Nose"},
+# //     {1,  "Neck"},
+# //     {2,  "RShoulder"},
+# //     {3,  "RElbow"},
+# //     {4,  "RWrist"},
+# //     {5,  "LShoulder"},
+# //     {6,  "LElbow"},
+# //     {7,  "LWrist"},
+# //     {8,  "MidHip"},
+# //     {9,  "RHip"},
+# //     {10, "RKnee"},
+# //     {11, "RAnkle"},
+# //     {12, "LHip"},
+# //     {13, "LKnee"},
+# //     {14, "LAnkle"},
+# //     {15, "REye"},
+# //     {16, "LEye"},
+# //     {17, "REar"},
+# //     {18, "LEar"},
+# //     {19, "LBigToe"},
+# //     {20, "LSmallToe"},
+# //     {21, "LHeel"},
+# //     {22, "RBigToe"},
+# //     {23, "RSmallToe"},
+# //     {24, "RHeel"},
+# //     {25, "Background"}
+
+
+class TRI3D_SmartBox:
+    
+    
+    def from_torch_image(self, image):
+        image = image.cpu().numpy() * 255.0
+        image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    def to_torch_image(self, image):
+        image = image.astype(dtype=np.float32)
+        image /= 255.0
+        image = torch.from_numpy(image)
+        return image
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE", ),
+                "keypoints_json": ("STRING", {"multiline": True}),
+            },
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE", )
+    CATEGORY = "TRI3D"
+
+    def extract_torso_keypoints(self, keypoints):
+        # Indices for torso-related keypoints
+        torso_indices = [8, 9, 10, 11, 12, 13]
+        return [keypoints[i] for i in torso_indices]
+
+
+    def run(self, image, keypoints_json):
+        
+        kp_data = json.loads(open(keypoints_json, 'r').read())
+        original_height, original_width = kp_data['height'], kp_data['width']
+        torso_keypoints = self.extract_torso_keypoints(kp_data['keypoints'])
+
+        # Convert Torch image to OpenCV format
+        cv_image = self.from_torch_image(image)
+        
+        # Remove the batch dimension if present
+        if len(cv_image.shape) == 4:
+            cv_image = cv_image[0]
+
+        # Adjust keypoints to match the image dimensions
+        adjusted_keypoints = self.adjust_keypoints(torso_keypoints, cv_image.shape, original_height, original_width)
+        
+
+        # Fill the area below the hip line
+        filled_image = self.fill_below_hip(cv_image, adjusted_keypoints)
+        
+
+        # Convert back to Torch format
+        torch_image = self.to_torch_image(filled_image)
+
+        # Add the batch dimension back
+        torch_image = torch_image.unsqueeze(0)
+        
+
+        return (torch_image,)
+
+    def adjust_keypoints(self, keypoints, image_shape, original_height, original_width):
+        image_height, image_width = image_shape[:2]
+        scale_x = image_width / original_width
+        scale_y = image_height / original_height
+
+        adjusted_keypoints = [
+            (int(x * scale_x), int(y * scale_y)) for x, y in keypoints
+        ]
+        return adjusted_keypoints
+
+    def fill_below_hip(self, image, keypoints):
+        # Correct the indices for hip keypoints
+        # Assuming indices 8 and 11 are for left and right hips
+        # print(keypoints,"hip keypoints")
+        try:
+            valid_y_coords = [kp[1] for kp in keypoints if kp[1] >= 0]
+            hip_y = min(valid_y_coords) if valid_y_coords else 0
+        except:
+            hip_y = 0
+        
+        if hip_y == 0:
+            return image
+
+        # Find the bounding box of the mask below the hip line
+        mask = image[:, :, 0]  # Assuming single-channel mask
+        below_hip = mask[hip_y:, :]
+        contours, _ = cv2.findContours(below_hip, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        cnt = 0
+        
+        for contour in contours:
+            x, y, w, h = cv2.boundingRect(contour)
+            # print(cnt, x,y,w,h, cv2.contourArea(contour), "cnt,x,y,w,h,area")
+            cnt+=1
+        #     cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (255, 255, 255), -1)
+        contours = [contour for contour in contours if cv2.contourArea(contour) > 20]
+
+        if len(contours) == 0:
+            return image
+        # Combine all contours into one
+        all_contours = np.vstack(contours)
+
+        # Calculate a single bounding rectangle for all contours
+        x, y, w, h = cv2.boundingRect(all_contours)
+        # print(x,y,w,h, "x,y,w,h")
+        cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (255, 255, 255), -1)
+
+        return image
+
+    
+class TRI3D_Skip_HeadMask:
+    def from_torch_image(self, image):
+        image = image.cpu().numpy() * 255.0
+        image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    def to_torch_image(self, image):
+        image = image.astype(dtype=np.float32)
+        image /= 255.0
+        image = torch.from_numpy(image)
+        return image
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE", ),
+                "head_mask": ("IMAGE", ),
+            },
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE", )
+    CATEGORY = "TRI3D"
+
+    def run(self, image, head_mask):
+        # Convert Torch images to OpenCV format
+        cv_image = self.from_torch_image(image)
+        cv_head_mask = self.from_torch_image(head_mask)
+
+        # Remove the batch dimension if present
+        if len(cv_image.shape) == 4:
+            cv_image = cv_image[0]
+        if len(cv_head_mask.shape) == 4:
+            cv_head_mask = cv_head_mask[0]
+
+        # Find the lowest point in the head mask
+        mask = cv_head_mask[:, :, 0]  # Assuming single-channel mask
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        lowest_y = 0
+        for contour in contours:
+            for point in contour:
+                x, y = point[0]
+                if y > lowest_y:
+                    lowest_y = y
+
+        # Black out everything above the lowest point
+        cv_image[:lowest_y, :] = 0
+
+        # Convert back to Torch format
+        torch_image = self.to_torch_image(cv_image)
+
+        # Add the batch dimension back
+        torch_image = torch_image.unsqueeze(0)
+
+        return (torch_image,)
+
+
+    
+class TRI3D_Skip_HeadMask_AddNeck:
+
+    def adjust_keypoints(self, keypoints, image_shape, original_height, original_width):
+        image_height, image_width = image_shape[:2]
+        scale_x = image_width / original_width
+        scale_y = image_height / original_height
+
+        adjusted_keypoints = [
+            (int(x * scale_x), int(y * scale_y)) for x, y in keypoints
+        ]
+        return adjusted_keypoints
+
+    def from_torch_image(self, image):
+        image = image.cpu().numpy() * 255.0
+        image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    def to_torch_image(self, image):
+        image = image.astype(dtype=np.float32)
+        image /= 255.0
+        image = torch.from_numpy(image)
+        return image
+
+
+    def extract_neck_keypoint(self, keypoints):
+        # Indices for torso-related keypoints
+        neck_indices = [1]
+        return [keypoints[i] for i in neck_indices]
+
+    def extract_ear_keypoints(self, keypoints):
+        # Indices for ear keypoints (17=right ear, 18=left ear)
+        ear_indices = [17, 18]
+        return [keypoints[i] for i in ear_indices]
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE", ),
+                "head_mask": ("IMAGE", ),
+                "keypoints_json": ("STRING", {"multiline": True}),
+                "ratio_aggression": ("FLOAT", {"default": 0.5, "min": 0, "max": 1, "step": 0.01}),
+                "neck_width_factor": ("FLOAT", {"default": 0.8, "min": 0.1, "max": 1.5, "step": 0.05}),
+            },
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE", )
+    CATEGORY = "TRI3D"
+
+    def run(self, image, head_mask, keypoints_json, ratio_aggression, neck_width_factor):
+        # Convert Torch images to OpenCV format
+        cv_image = self.from_torch_image(image)
+        cv_head_mask = self.from_torch_image(head_mask)
+
+        # Remove the batch dimension if present
+        if len(cv_image.shape) == 4:
+            cv_image = cv_image[0]
+        if len(cv_head_mask.shape) == 4:
+            cv_head_mask = cv_head_mask[0]
+
+        kp_data = json.loads(open(keypoints_json, 'r').read())
+        original_height, original_width = kp_data['height'], kp_data['width']
+        neck_keypoints = self.extract_neck_keypoint(kp_data['keypoints'])
+        
+        # Make a copy of the original image
+        result_image = cv_image.copy()
+        
+        # Adjust keypoints to match the image dimensions
+        adjusted_neck_keypoints = self.adjust_keypoints(neck_keypoints, cv_image.shape, original_height, original_width)
+        
+        # Find the lowest point and face dimensions in the head mask
+        mask = cv_head_mask[:, :, 0]  # Assuming single-channel mask
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+        # Find the chin point (lowest point) and calculate face properties
+        lowest_y = 0
+        face_center_x = cv_image.shape[1] // 2  # Default to center of image
+        face_width = cv_image.shape[1] // 3     # Default face width
+        
+        if contours:
+            # Find the lowest point (chin)
+            for contour in contours:
+                for point in contour:
+                    x, y = point[0]
+                    if y > lowest_y:
+                        lowest_y = y
+            
+            # Calculate face bounding box and center of gravity
+            x, y, w, h = cv2.boundingRect(contours[0])
+            face_width = w
+            
+            # Calculate center of gravity of the face mask
+            M = cv2.moments(contours[0])
+            if M["m00"] != 0:
+                face_center_x = int(M["m10"] / M["m00"])
+            else:
+                face_center_x = x + w // 2
+        
+        # Calculate weighted average point between neck and chin
+        neck_y = adjusted_neck_keypoints[0][1]
+        if neck_y <= 0:
+            neck_y = lowest_y
+        average_y = int((neck_y * ratio_aggression + lowest_y * (1 - ratio_aggression)))
+        
+        print(neck_y, lowest_y, "neck_y, lowest_y")
+        print(average_y, "average_y")
+        
+        # ZONE 1: Black out everything above the chin point
+        result_image[:lowest_y, :] = 0
+        
+        # ZONE 2: Create a triangle for the neck area
+        if lowest_y < average_y:  # Only process if there's a gap between chin and average_y
+            # Create a mask for Zone 2
+            zone2_mask = np.zeros_like(cv_image[:,:,0])
+            
+            # Create a triangle with apex at weighted average point and base at chin level
+            # Apply the neck width factor to the face width
+            neck_width = int(face_width * neck_width_factor)
+            triangle_half_width = neck_width // 2
+            
+            # Create polygon points for the triangle
+            triangle_points = np.array([
+                [face_center_x, average_y],  # Apex at weighted average point
+                [face_center_x - triangle_half_width, lowest_y],  # Left base point at chin level
+                [face_center_x + triangle_half_width, lowest_y]   # Right base point at chin level
+            ], dtype=np.int32)
+            
+            # Fill the triangle in the mask
+            cv2.fillPoly(zone2_mask, [triangle_points], 255)
+            
+            # Apply the mask only to the region between chin and weighted average
+            for y in range(lowest_y, average_y):
+                for x in range(cv_image.shape[1]):
+                    if zone2_mask[y, x] > 0:
+                        result_image[y, x] = 0
+        
+        # ZONE 3: Area below weighted average point is left as is
+        # No action needed for this zone
+        
+        # Convert back to Torch format
+        torch_image = self.to_torch_image(result_image)
+
+        # Add the batch dimension back
+        torch_image = torch_image.unsqueeze(0)
+
+        return (torch_image,)
+
+
+class TRI3D_Image_extend:
+    def from_torch_image(self, image):
+        image = image.cpu().numpy() * 255.0
+        image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    def to_torch_image(self, image):
+        image = image.astype(dtype=np.float32)
+        image /= 255.0
+        image = torch.from_numpy(image)
+        return image
+    
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "face_mask": ("IMAGE", ),
+                "image": ("IMAGE", ),
+                "ratio": ("FLOAT", {"default": 1.5, "min": 1.2, "max": 2, "step": 0.01}),
+            },
+        }
+    
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE", "IMAGE", )
+    RETURN_NAMES = ("image", "mask_image", )
+    CATEGORY = "TRI3D"
+
+    def run(self, face_mask, image, ratio):
+        cv_face_mask = self.from_torch_image(face_mask)
+        cv_image = self.from_torch_image(image)
+
+        # Remove the batch dimension if present
+        if len(cv_image.shape) == 4:
+            cv_image = cv_image[0]
+        if len(cv_face_mask.shape) == 4:
+            cv_face_mask = cv_face_mask[0]
+        mask = cv_face_mask[:, :, 0]  # Assuming single-channel mask
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        lowest_y = 0
+        highest_y = cv_image.shape[0]
+        for contour in contours:
+            for point in contour:
+                x, y = point[0]
+                if y > lowest_y:
+                    lowest_y = y
+                if y < highest_y:
+                    highest_y = y
+
+        y_below_face = cv_image.shape[0] - lowest_y
+        y_face = lowest_y-highest_y
+
+        # Only extend if the space below face is less than 1.5 times face height
+        target_below_face = int(y_face * ratio)
+        # print("y_face", y_face)
+        # print("lowest_y", lowest_y)
+        # print("highest_y", highest_y)
+        # print("target_below_face", target_below_face)
+        # print("y_below_face", y_below_face)
+
+        original_height = cv_image.shape[0]
+        original_width = cv_image.shape[1]
+        
+        if y_below_face < target_below_face:
+            y_extend = target_below_face - y_below_face
+            
+            # Calculate how much to extend horizontally to maintain aspect ratio
+            new_height = original_height + y_extend
+            new_width = int(original_width * (new_height / original_height))
+            x_extend = new_width - original_width
+            x_extend_left = x_extend // 2
+            x_extend_right = x_extend - x_extend_left
+            
+            # Extend the image in all necessary directions
+            cv_image = cv2.copyMakeBorder(
+                cv_image, 
+                0, y_extend,                    # top, bottom
+                x_extend_left, x_extend_right,  # left, right
+                cv2.BORDER_CONSTANT, 
+                value=[0, 0, 0]
+            )
+            
+            # Create extension mask
+            extension_mask = np.zeros_like(cv_image)
+            # Make extended portions white
+            extension_mask[original_height:, :] = 255  # bottom extension
+            extension_mask[:, :x_extend_left] = 255    # left extension
+            extension_mask[:, -x_extend_right:] = 255  # right extension
+
+        else:
+            extension_mask = np.zeros_like(cv_image)
+
+        # Convert both images back to torch format
+        torch_image = self.to_torch_image(cv_image)
+        torch_mask = self.to_torch_image(extension_mask)
+        
+        # Add batch dimension to both
+        torch_image = torch_image.unsqueeze(0)
+        torch_mask = torch_mask.unsqueeze(0)
+        
+        return (torch_image, torch_mask)
+
+
+class TRI3D_Smart_Depth:
+    
+    
+    def from_torch_image(self, image):
+        image = image.cpu().numpy() * 255.0
+        image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    def to_torch_image(self, image):
+        image = image.astype(dtype=np.float32)
+        image /= 255.0
+        image = torch.from_numpy(image)
+        return image
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE", ),
+                "keypoints_json": ("STRING", {"multiline": True}),
+            },
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE", )
+    CATEGORY = "TRI3D"
+
+    def extract_torso_keypoints(self, keypoints):
+        # Indices for torso-related keypoints
+        torso_indices = [8, 9, 10, 11, 12, 13]
+        return [keypoints[i] for i in torso_indices]
+
+
+    def run(self, image, keypoints_json):
+        
+        kp_data = json.loads(open(keypoints_json, 'r').read())
+        original_height, original_width = kp_data['height'], kp_data['width']
+        torso_keypoints = self.extract_torso_keypoints(kp_data['keypoints'])
+
+        # Convert Torch image to OpenCV format
+        cv_image = self.from_torch_image(image)
+        
+        # Remove the batch dimension if present
+        if len(cv_image.shape) == 4:
+            cv_image = cv_image[0]
+
+        # Adjust keypoints to match the image dimensions
+        adjusted_keypoints = self.adjust_keypoints(torso_keypoints, cv_image.shape, original_height, original_width)
+        
+
+        # Fill the area below the hip line
+        filled_image = self.fill_below_hip(cv_image, adjusted_keypoints)
+        
+
+        # Convert back to Torch format
+        torch_image = self.to_torch_image(filled_image)
+
+        # Add the batch dimension back
+        torch_image = torch_image.unsqueeze(0)
+        
+
+        return (torch_image,)
+
+    def adjust_keypoints(self, keypoints, image_shape, original_height, original_width):
+        image_height, image_width = image_shape[:2]
+        scale_x = image_width / original_width
+        scale_y = image_height / original_height
+
+        adjusted_keypoints = [
+            (int(x * scale_x), int(y * scale_y)) for x, y in keypoints
+        ]
+        return adjusted_keypoints
+
+    def fill_below_hip(self, image, keypoints):
+        # Correct the indices for hip keypoints
+        # Assuming indices 8 and 11 are for left and right hips
+        # print(keypoints,"hip keypoints")
+        try:
+            valid_y_coords = [kp[1] for kp in keypoints if kp[1] >= 0]
+            hip_y = min(valid_y_coords) if valid_y_coords else 0
+        except:
+            hip_y = 0
+        
+        if hip_y == 0:
+            return image
+
+        # Find the bounding box of the mask below the hip line
+        mask = image[:, :, 0]  # Assuming single-channel mask
+        below_hip = mask[hip_y:, :]
+        contours, _ = cv2.findContours(below_hip, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        cnt = 0
+        
+        for contour in contours:
+            x, y, w, h = cv2.boundingRect(contour)
+            # print(cnt, x,y,w,h, cv2.contourArea(contour), "cnt,x,y,w,h,area")
+            cnt+=1
+        #     cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (255, 255, 255), -1)
+        contours = [contour for contour in contours if cv2.contourArea(contour) > 0]
+
+        if len(contours) == 0:
+            return image
+        # Combine all contours into one
+        all_contours = np.vstack(contours)
+
+        # Calculate a single bounding rectangle for all contours
+        x, y, w, h = cv2.boundingRect(all_contours)
+        # print(x,y,w,h, "x,y,w,h")
+        cv2.rectangle(image, (x, y + hip_y), (x + w, y + h + hip_y), (0, 0, 0), -1)
+
+        return image
+
+
+class TRI3D_NarrowfyImage:
+    def from_torch_image(self, image):
+        image = image.cpu().numpy() * 255.0
+        image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    def to_torch_image(self, image):
+        image = image.astype(dtype=np.float32)
+        image /= 255.0
+        image = torch.from_numpy(image)
+        return image
+    
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE", ),
+                "mask": ("IMAGE", ),
+                "aspect_ratio": ("FLOAT", {"default": 0.33, "min": 0.25, "max": 1, "step": 0.01}),
+                "border_margin": ("INT", {"default": 15, "min": 10, "max": 100, "step": 1}),
+            },
+        }
+    
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE", "IMAGE", "INT", "INT",)
+    RETURN_NAMES = ("cropped_image", "cropped_mask", "cropped_width", "cropped_height",)
+    CATEGORY = "TRI3D"
+
+    def run(self, image, mask, aspect_ratio, border_margin):
+        # Convert to CV format and remove batch dimension
+        cv_image = self.from_torch_image(image)[0]
+        cv_mask = self.from_torch_image(mask)[0]
+        
+        # Find bounding box of the mask
+        mask_channel = cv_mask[:, :, 0]
+        contours, _ = cv2.findContours(mask_channel, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+        if not contours:
+            return image, mask, aspect_ratio
+            
+        # Filter contours by area
+        significant_contours = [cnt for cnt in contours if cv2.contourArea(cnt) > 100]
+        
+        if not significant_contours:
+            return image, mask, aspect_ratio
+            
+        # Get combined bounding box for all significant contours
+        x_min = float('inf')
+        y_min = float('inf')
+        x_max = 0
+        y_max = 0
+        
+        for contour in significant_contours:
+            x, y, w, h = cv2.boundingRect(contour)
+            x_min = min(x_min, x)
+            y_min = min(y_min, y)
+            x_max = max(x_max, x + w)
+            y_max = max(y_max, y + h)
+        
+        # Calculate final width and height with margin
+        margin = border_margin
+        x = max(0, x_min - margin)  # Ensure we don't go below 0
+        y = max(0, y_min - margin)
+        w = min(cv_image.shape[1] - x, (x_max - x_min) + 2 * margin)  # Ensure we don't exceed image width
+        h = min(cv_image.shape[0] - y, (y_max - y_min) + 2 * margin)  # Ensure we don't exceed image height
+        
+        # Crop both image and mask to bounding box
+        cropped_image = cv_image[y:y+h, x:x+w]
+        cropped_mask = cv_mask[y:y+h, x:x+w]
+        
+        # Calculate required height for aspect ratio 1/3
+        min_height = w * 1/aspect_ratio
+        if h < min_height:
+            height_extend = min_height - h
+            
+            # Extend image with black pixels
+            extended_image = cv2.copyMakeBorder(
+                cropped_image,
+                0, int(height_extend),  # top, bottom
+                0, 0,                   # left, right
+                cv2.BORDER_CONSTANT,
+                value=[0, 0, 0]
+            )
+            
+            # Create mask with white pixels only in extended region
+            extended_mask = cv2.copyMakeBorder(
+                np.zeros_like(cropped_mask),  # Start with black base
+                0, int(height_extend),  # top, bottom
+                0, 0,                   # left, right
+                cv2.BORDER_CONSTANT,
+                value=[255, 255, 255]   # White extension
+            )
+            
+            cropped_image = extended_image
+            cropped_mask = extended_mask
+        
+        # Convert back to torch format and add batch dimension
+        torch_image = self.to_torch_image(cropped_image).unsqueeze(0)
+        torch_mask = self.to_torch_image(cropped_mask).unsqueeze(0)
+        
+        return (torch_image, torch_mask,w,h)
+
+
+
+
+class TRI3D_CropAndExtend:
+    def from_torch_image(self, image):
+        image = image.cpu().numpy() * 255.0
+        image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    def to_torch_image(self, image):
+        image = image.astype(dtype=np.float32)
+        image /= 255.0
+        image = torch.from_numpy(image)
+        return image
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "garment_image": ("IMAGE",),
+                "garment_mask": ("IMAGE",),
+                "human_image": ("IMAGE",),
+                "human_mask": ("IMAGE",),
+                "margin": ("INT", {"default": 10, "min": 0, "max": 50}),
+            },
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE", "IMAGE", "INT", "INT",)
+    RETURN_NAMES = ("cropped_garment", "cropped_garment_mask", "cropped_human", "cropped_human_mask", "cropped_width", "cropped_height",)
+    
+    def run(self, garment_image, garment_mask, human_image, human_mask, margin):
+        # Convert to CV format and remove batch dimension
+        cv_garment = self.from_torch_image(garment_image)[0]
+        cv_garment_mask = self.from_torch_image(garment_mask)[0]
+        cv_human = self.from_torch_image(human_image)[0]
+        cv_human_mask = self.from_torch_image(human_mask)[0]
+        
+        # Process garment
+        mask_channel = cv_garment_mask[:, :, 0]
+        contours, _ = cv2.findContours(mask_channel, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+        if not contours:
+            return garment_image, garment_mask, human_image, human_mask, cv_garment.shape[1], cv_garment.shape[0]
+            
+        # Get bounding box with margin
+        x, y, w, h = cv2.boundingRect(contours[0])
+        x = max(0, x - margin)
+        y = max(0, y - margin)
+        w = min(cv_garment.shape[1] - x, w + 2 * margin)
+        h = min(cv_garment.shape[0] - y, h + 2 * margin)
+        
+        # Store the cropped dimensions before extension
+        cropped_width = w
+        cropped_height = h
+        
+        # Crop garment and its mask
+        cropped_garment = cv_garment[y:y+h, x:x+w]
+        cropped_garment_mask = cv_garment_mask[y:y+h, x:x+w]
+        
+        # Calculate required height for aspect ratio 1/3
+        min_height = w * 3
+        if h < min_height:
+            height_extend = min_height - h
+            
+            # Extend garment image and mask
+            extended_garment = cv2.copyMakeBorder(
+                cropped_garment,
+                0, int(height_extend),
+                0, 0,
+                cv2.BORDER_CONSTANT,
+                value=[0, 0, 0]
+            )
+            
+            extended_garment_mask = cv2.copyMakeBorder(
+                cropped_garment_mask,
+                0, int(height_extend),
+                0, 0,
+                cv2.BORDER_CONSTANT,
+                value=[255, 255, 255]
+            )
+            
+            cropped_garment = extended_garment
+            cropped_garment_mask = extended_garment_mask
+        
+        # Process human image similarly
+        mask_channel = cv_human_mask[:, :, 0]
+        contours, _ = cv2.findContours(mask_channel, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+        if contours:
+            x, y, w, h = cv2.boundingRect(contours[0])
+            x = max(0, x - margin)
+            y = max(0, y - margin)
+            w = min(cv_human.shape[1] - x, w + 2 * margin)
+            h = min(cv_human.shape[0] - y, h + 2 * margin)
+            
+            cropped_human = cv_human[y:y+h, x:x+w]
+            cropped_human_mask = cv_human_mask[y:y+h, x:x+w]
+            
+            min_height = w * 3
+            if h < min_height:
+                height_extend = min_height - h
+                
+                extended_human = cv2.copyMakeBorder(
+                    cropped_human,
+                    0, int(height_extend),
+                    0, 0,
+                    cv2.BORDER_CONSTANT,
+                    value=[0, 0, 0]
+                )
+                
+                extended_human_mask = cv2.copyMakeBorder(
+                    cropped_human_mask,
+                    0, int(height_extend),
+                    0, 0,
+                    cv2.BORDER_CONSTANT,
+                    value=[255, 255, 255]
+                )
+                
+                cropped_human = extended_human
+                cropped_human_mask = extended_human_mask
+        
+        # Convert back to torch format and add batch dimension
+        torch_garment = self.to_torch_image(cropped_garment).unsqueeze(0)
+        torch_garment_mask = self.to_torch_image(cropped_garment_mask).unsqueeze(0)
+        torch_human = self.to_torch_image(cropped_human).unsqueeze(0)
+        torch_human_mask = self.to_torch_image(cropped_human_mask).unsqueeze(0)
+        
+        return (torch_garment, torch_garment_mask, torch_human, torch_human_mask, cropped_width, cropped_height)
+
+class TRI3D_Skip_LipMask:
+
+    def adjust_keypoints(self, keypoints, image_shape, original_height, original_width):
+        image_height, image_width = image_shape[:2]
+        scale_x = image_width / original_width
+        scale_y = image_height / original_height
+
+        adjusted_keypoints = [
+            (int(x * scale_x), int(y * scale_y)) for x, y in keypoints
+        ]
+        return adjusted_keypoints
+
+    def from_torch_image(self, image):
+        image = image.cpu().numpy() * 255.0
+        image = np.clip(image, 0, 255).astype(np.uint8)
+        return image
+
+    def to_torch_image(self, image):
+        image = image.astype(dtype=np.float32)
+        image /= 255.0
+        image = torch.from_numpy(image)
+        return image
+
+    def extract_lip_keypoints(self, keypoints):
+        # In DWPose, lips are typically keypoints in face area
+        # Assuming standard face keypoint format where lips are around indices 61-68
+        # This may need adjustment based on your specific keypoint format
+        lip_indices = range(61, 69)  # Adjust these indices based on your keypoint format
+        
+        # Filter out invalid keypoints (those with negative confidence or coordinates)
+        lip_keypoints = []
+        for idx in lip_indices:
+            if idx < len(keypoints):
+                x, y = keypoints[idx]
+                if x >= 0 and y >= 0:  # Check for valid coordinates
+                    lip_keypoints.append((x, y))
+        
+        return lip_keypoints
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE", ),
+                "keypoints_json": ("STRING", {"multiline": True}),
+            },
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE", )
+    CATEGORY = "TRI3D"
+
+    def run(self, image, keypoints_json):
+        # Convert Torch image to OpenCV format
+        cv_image = self.from_torch_image(image)
+
+        # Remove the batch dimension if present
+        if len(cv_image.shape) == 4:
+            cv_image = cv_image[0]
+
+        # Make a copy of the original image
+        result_image = cv_image.copy()
+        
+        # Parse keypoints JSON
+        try:
+            kp_data = json.loads(open(keypoints_json, 'r').read())
+            original_height, original_width = kp_data['height'], kp_data['width']
+            keypoints = kp_data['keypoints']
+            
+            # Extract lip keypoints
+            lip_keypoints = self.extract_lip_keypoints(keypoints)
+            
+            # If no valid lip keypoints found, use a fallback approach
+            if not lip_keypoints:
+                # Fallback: use the nose point (index 0) as reference
+                nose_point = keypoints[0]
+                if nose_point[1] > 0:  # If y-coordinate is valid
+                    # Estimate lip position slightly below nose
+                    lip_y = int(nose_point[1] + 0.15 * cv_image.shape[0])
+                    lowest_y = lip_y
+                else:
+                    # If no valid reference point, use 1/3 of the image height
+                    lowest_y = cv_image.shape[0] // 3
+            else:
+                # Find the lowest y-coordinate among lip keypoints
+                adjusted_lip_keypoints = self.adjust_keypoints(lip_keypoints, cv_image.shape, original_height, original_width)
+                lowest_y = max([kp[1] for kp in adjusted_lip_keypoints])
+            
+            # Black out everything above the lowest lip point
+            result_image[:lowest_y, :] = 0
+            
+        except Exception as e:
+            print(f"Error processing keypoints JSON: {e}")
+            # In case of error, return the original image
+            result_image = cv_image
+        
+        # Convert back to Torch format
+        torch_image = self.to_torch_image(result_image)
+
+        # Add the batch dimension back
+        torch_image = torch_image.unsqueeze(0)
+
+        return (torch_image,)

string_check.py

@@ -0,0 +1,42 @@
+import os
+
+class TRI3D_StringContains:
+    """
+    ComfyUI node that checks if a specified string exists within another string.
+    Performs case-insensitive comparison by converting all text to lowercase.
+    """
+    
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "input_string": ("STRING", {"multiline": True}),
+                "search_string": ("STRING", {"default": "", "multiline": False}),
+            },
+        }
+    
+    FUNCTION = "run"
+    RETURN_TYPES = ("BOOLEAN",)
+    CATEGORY = "TRI3D"
+    
+    def run(self, input_string, search_string):
+        # Convert both strings to lowercase for case-insensitive comparison
+        input_lower = input_string.lower()
+        search_lower = search_string.lower()
+        
+        # Check if search string exists in input string
+        contains = search_lower in input_lower
+        
+        return (contains,)
+
+# # Node registration for ComfyUI
+# NODE_CLASS_MAPPINGS = {
+#     "TRI3D_StringContains": TRI3D_StringContains
+# }
+
+# NODE_DISPLAY_NAME_MAPPINGS = {
+#     "TRI3D_StringContains": "TRI3D String Contains"
+# }

utility_nodes.py

@@ -0,0 +1,409 @@
+import torch, cv2, json
+import numpy as np
+
+
+def from_torch_image(image):
+    image = image.cpu().numpy() * 255.0
+    image = np.clip(image, 0, 255).astype(np.uint8)
+    return image
+
+
+def to_torch_image(image):
+    image = image.astype(dtype=np.float32)
+    image /= 255.0
+    image = torch.from_numpy(image)
+    return image
+
+class TRI3D_clean_mask():
+
+    """For the given mask and threshold area, remove all patches in the mask with area smaller than threshold"""
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "masks": ("MASK", ),
+                "threshold":("FLOAT",{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01})
+            }
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("MASK", "BOOL")
+    RETURN_NAMES = ("mask", "cleaned")
+    CATEGORY = "TRI3D"
+
+    def run(self, masks, threshold):
+        batch_results = []
+        for mask in masks:
+            mask = from_torch_image(mask)
+            mask = np.where(mask < 127, 0, 255).astype(np.uint8)
+            h,w = mask.shape[:2]
+
+            total_area = h*w
+            # num_labels, labels = cv2.connectedComponents(mask)
+            region_mask = np.zeros_like(mask)
+
+            # for label in range(1, num_labels):  
+            #     area_percent = (np.sum(labels == label)/ total_area) * 100
+            #     if area_percent < threshold:
+            #         continue
+            #     region_mask[labels == label] = 255
+            less_than_threshold = True
+            area_percent = (np.sum(mask == 255)/ total_area) * 100
+            if area_percent > threshold:
+                region_mask[mask == 255] = 255
+                less_than_threshold = False
+            region_mask = to_torch_image(region_mask)
+            batch_results.append(region_mask.squeeze(0))
+
+        batch_results = torch.stack(batch_results)
+        return (batch_results, less_than_threshold)
+        
+
+class TRI3D_extract_pose_part():
+    """
+        For the given pose, extract region around body parts, region can be defined by % of image size  
+    """
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE", ),
+                "pose_json": ("STRING",{"default" : "dwpose/keypoints/input.json"}),
+                "width_pad": ("FLOAT",{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
+                "height_pad": ("FLOAT",{"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}),
+                "shoulders":("BOOLEAN", {
+                    "default": False
+                })
+            }
+        }
+
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE", "STRING")
+    RETURN_NAMES = ("image", "coords")
+    CATEGORY = "TRI3D"
+
+    def get_frame_coords(self,point1, point2):
+        x1, y1 = point1
+        x2, y2 = point2
+
+        xmin, xmax, ymin, ymax = min(x1, x2), max(x1, x2), min(y1, y2), max(y1, y2)
+
+        for i in [xmin, xmax, ymin, ymax]:
+            if i < 0: 
+                return None
+
+        return [xmin, xmax, ymin, ymax]
+
+    def run(self, image, pose_json, width_pad, height_pad, shoulders):
+        """
+        image : input image
+        width_pad: % of image width you want to apply on both size of pose body part
+        height_pad: % of image width you want to apply on both size of pose body part
+        rest of them are body parts
+        """
+
+        image = from_torch_image(image[0])
+        batch_result = []
+        input_pose = json.load(open(pose_json))
+        keypoints = input_pose['keypoints']
+
+        og_h, og_w = image.shape[:2]
+        ph, pw = [input_pose['height'], input_pose['width']]
+
+        for i,point in enumerate(keypoints):
+            x,y = point
+            y = int((y/ph)*og_h)
+            x = int((x/pw)*og_w)
+            keypoints[i] = [x, y]
+
+        width_offset = int(og_w * (width_pad) / 100)
+        height_offset = int(og_h * (height_pad) / 100)
+
+        xmin, xmax, ymin, ymax = [0, og_w, 0, og_h]
+
+        part_to_coords = {
+            "shoulders":self.get_frame_coords(keypoints[2], keypoints[5])
+        }
+
+        if shoulders:
+            print(part_to_coords["shoulders"])
+            if part_to_coords["shoulders"] != None:
+                new_xmin, new_xmax, new_ymin, new_ymax = part_to_coords["shoulders"]
+
+                xmin, xmax, ymin, ymax = new_xmin, new_xmax, new_ymin, new_ymax
+        
+        xmin = max(0, xmin - width_offset)
+        xmax = min(og_w, xmax + width_offset)
+        ymin = max(0, ymin - height_offset)
+        ymax = min(og_h, ymax + height_offset)
+
+        image = image[ymin:ymax, xmin:xmax, :].astype(np.uint8)
+        image = to_torch_image(image)
+        batch_result.append(image)
+        batch_result = torch.stack(batch_result)
+        print("final_coords", xmin, xmax, ymin, ymax)
+        coords = ",".join([str(xmin), str(xmax), str(ymin), str(ymax)])
+        
+        return batch_result, coords
+
+class TRI3D_position_pose_part():
+    """
+        put back extracted parts on OG image 
+    """
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "og_image": ("IMAGE", ),
+                "extracted_image": ("IMAGE", ),
+                "coords": ("STRING",{"default" : "xmin, xmax, ymin, ymax"}),
+            }
+        }
+
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE", )
+    RETURN_NAMES = ("image", )
+    CATEGORY = "TRI3D"
+
+    def run(self, og_image, extracted_image, coords):
+
+        batch_result = []
+        og_image = from_torch_image(og_image[0])
+        extracted_image = from_torch_image(extracted_image[0])
+        
+        xmin, xmax, ymin, ymax = [int(i) for i in coords.split(",")]
+
+        og_image[ymin:ymax, xmin:xmax, :] = extracted_image
+
+        og_image = to_torch_image(og_image).unsqueeze(0)
+        batch_result.append(og_image)
+        batch_result = torch.stack(batch_result)
+        return batch_result
+
+
+class TRI3D_fill_mask():
+    """
+        fill mask with the neighbouring pixels
+    """
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE", ),
+                "mask": ("MASK", ),
+                "negative_mask": ("MASK", ),
+                "offset":("FLOAT",{"default": 1, "min": 0.0, "max": 100.0, "step": 0.01})
+            }
+        }
+
+    FUNCTION = "run"
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    CATEGORY = "TRI3D"
+
+    def run(self, image, mask, negative_mask, offset):
+        image = from_torch_image(image[0])
+
+        mask = mask[0].cpu().numpy()
+        mask = np.expand_dims(mask, -1)
+        mh, mw, _ = mask.shape
+
+        inverse_mask = np.ones_like(mask) - mask
+        
+        negative_mask = negative_mask[0].cpu().numpy()
+        indices = np.where(mask > 0)
+
+        offset = offset / 100
+        
+        source = image.copy()
+
+        for y,x in zip(indices[0],indices[1]):
+            x_off = min(mw-1, int(x + offset * mw))
+            if negative_mask[y][x_off] == 0:         #check if pixles on right are outside body
+                source[y][x] = image[y][x_off]
+                
+            else:
+                x_off = max(0, int(x - offset * mw))       #check if pixles on left are outside body
+                if negative_mask[y][x_off] == 0:
+                    source[y][x] = image[y][x_off]
+                else:
+                    y_off = max(0, int(y - offset * mh))
+                    if negative_mask[y_off][x] == 0:            #check if pixles on top are outside body
+                        source[y][x] = image[y_off][x]
+
+                    else:
+                        y_off = min(mh-1, int(y + offset * mh))
+                        if negative_mask[y_off][x] == 0:            #check if pixles on bottom are outside body
+                            source[y][x] = image[y_off][x]
+
+        image = mask * source + inverse_mask * image
+        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,)
+
+class TRI3D_extract_facer_mask:
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "background": ("BOOLEAN", {
+                    "default": False
+                }),
+                'hair':("BOOLEAN", {
+                    "default": False
+                }),
+                'lower_lip':("BOOLEAN", {
+                    "default": False
+                }),
+                'inner_mouth':("BOOLEAN", {
+                    "default": False
+                }),
+                'upper_lip':("BOOLEAN", {
+                    "default": False
+                }),
+                'nose':("BOOLEAN", {
+                    "default": False
+                }),
+                'left_eyebrow':("BOOLEAN", {
+                    "default": False
+                }),
+                'right_eyebrow':("BOOLEAN", {
+                    "default": False
+                }),
+                'left_eye':("BOOLEAN", {
+                    "default": False
+                }),
+                'right_eye':("BOOLEAN", {
+                    "default": False
+                }),
+                'face':("BOOLEAN", {
+                    "default": False
+                })
+            }
+        }
+
+    RETURN_TYPES = ("MASK", )
+    FUNCTION = "main"
+    CATEGORY = "TRI3D"
+
+    def main(self, image, background, hair, lower_lip, inner_mouth, upper_lip, nose, left_eyebrow, right_eyebrow, left_eye, right_eye, face):
+
+        image = from_torch_image(image[0])
+        h,w,_ = image.shape
+        
+        mask = np.zeros_like(image)
+        
+        label_to_rgb = {'background':[0,0,0], 'face':[0,138,255], 'right_eye':[180, 255, 0], 'left_eye':[42, 255, 0], 'right_eyebrow':[0, 255, 96],
+                'left_eyebrow':[0,255,234], 'nose':[255, 192, 0],  'upper_lip':[255, 54, 0], 'inner_mouth':[255, 0, 84], 'lower_lip':[255, 0, 222], 
+                'hair':[150,0,255]}
+
+        if background:
+            temp = np.all(image == label_to_rgb['background'], axis=-1)
+            idcs = np.where(temp==True)
+            mask[idcs] = 255
+        
+        if face:
+            temp = np.all(image == label_to_rgb['face'], axis=-1)
+            idcs = np.where(temp==True)
+            mask[idcs] = 255
+
+        if right_eye:
+            temp = np.all(image == label_to_rgb['right_eye'], axis=-1)
+            idcs = np.where(temp==True)
+            mask[idcs] = 255
+
+        if left_eye:
+            temp = np.all(image == label_to_rgb['left_eye'], axis=-1)
+            idcs = np.where(temp==True)
+            mask[idcs] = 255
+
+        if right_eyebrow:
+            temp = np.all(image == label_to_rgb['right_eyebrow'], axis=-1)
+            idcs = np.where(temp==True)
+            mask[idcs] = 255
+
+        if left_eyebrow:
+            temp = np.all(image == label_to_rgb['left_eyebrow'], axis=-1)
+            idcs = np.where(temp==True)
+            mask[idcs] = 255
+
+        if nose:
+            temp = np.all(image == label_to_rgb['nose'], axis=-1)
+            idcs = np.where(temp==True)
+            mask[idcs] = 255
+
+        if upper_lip:
+            temp = np.all(image == label_to_rgb['upper_lip'], axis=-1)
+            idcs = np.where(temp==True)
+            mask[idcs] = 255
+
+        if inner_mouth:
+            temp = np.all(image == label_to_rgb['inner_mouth'], axis=-1)
+            idcs = np.where(temp==True)
+            mask[idcs] = 255
+
+        if lower_lip:
+            temp = np.all(image == label_to_rgb['lower_lip'], axis=-1)
+            idcs = np.where(temp==True)
+            mask[idcs] = 255
+
+        if hair:
+            temp = np.all(image == label_to_rgb['hair'], axis=-1)
+            idcs = np.where(temp==True)
+            mask[idcs] = 255
+
+        mask = to_torch_image(mask[:,:,0]).unsqueeze(0)   
+        return (mask,)