Merge pull request #7 from Layer-norm/dense126

Dense126

README.md

@@ -36,6 +36,7 @@ All credit & copyright goes to https://github.com/lllyasviel.
 * Added YOLO-NAS models which are drop-in replacements of YOLOX
 * Fixed Openpose Face/Hands no longer detecting: https://github.com/Fannovel16/comfyui_controlnet_aux/issues/54
 * Added TorchScript implementation of DWPose and AnimalPose
+* Added TorchScript implementation of DensePose from [Colab notebook](https://colab.research.google.com/drive/16hcaaKs210ivpxjoyGNuvEXZD4eqOOSQ) which doesn't require detectron2. [Example](#densepose). Ps/r: Currently doesn't work
 # Q&A:
 ## Why some nodes doesn't appear after I installed this repo?
 
@@ -213,6 +214,9 @@ Credit to https://huggingface.co/thibaud/controlnet-sd21. You can get the same k
 ### Animal Pose (AP-10K)
 ![](./example_animal_pose.png)
 
+### DensePose
+![](./example_densepose.png)
+
 ## Semantic Segmantation
 ### OneFormer ADE20K Segmentor
 ![](https://huggingface.co/thibaud/controlnet-sd21/resolve/main/example_ade20k.png)

node_wrappers/densepose.py

@@ -0,0 +1,31 @@
+from ..utils import common_annotator_call, annotator_ckpts_path, HF_MODEL_NAME, create_node_input_types
+import comfy.model_management as model_management
+
+class DensePose_Preprocessor:
+    @classmethod
+    def INPUT_TYPES(s):
+        return create_node_input_types(
+            model=(["densepose_r50_fpn_dl.torchscript", "densepose_r101_fpn_dl.torchscript"], {"default": "densepose_r50_fpn_dl.torchscript"}),
+            cmap=(["Viridis (MagicAnimate)", "Parula (CivitAI)"], {"default": "Viridis (MagicAnimate)"})
+        )
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "execute"
+
+    CATEGORY = "ControlNet Preprocessors/Faces and Poses"
+
+    def execute(self, image, model, cmap, resolution=512):
+        from controlnet_aux.densepose import DenseposeDetector
+        return (common_annotator_call(
+                DenseposeDetector.from_pretrained("LayerNorm/DensePose-TorchScript-with-hint-image", model).to(model_management.get_torch_device()), 
+                image, 
+                cmap="viridis" if "Viridis" in cmap else "parula", 
+                resolution=resolution), )
+
+
+NODE_CLASS_MAPPINGS = {
+    "DensePosePreprocessor": DensePose_Preprocessor
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "DensePosePreprocessor": "DensePose Estimation"
+}

src/controlnet_aux/densepose/__init__.py

@@ -0,0 +1,68 @@
+import torchvision # Fix issue Unknown builtin op: torchvision::nms
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+from einops import rearrange
+from PIL import Image
+
+from controlnet_aux.util import HWC3, resize_image_with_pad, common_input_validate, annotator_ckpts_path, custom_hf_download
+from .densepose import DensePoseMaskedColormapResultsVisualizer, _extract_i_from_iuvarr, densepose_chart_predictor_output_to_result_with_confidences
+
+N_PART_LABELS = 24
+
+class DenseposeDetector:
+    def __init__(self, model):
+        self.dense_pose_estimation = model
+        self.device = "cpu"
+        self.result_visualizer = DensePoseMaskedColormapResultsVisualizer(
+            cmap=cv2.COLORMAP_PARULA, 
+            alpha=1, 
+            data_extractor=_extract_i_from_iuvarr, 
+            segm_extractor=_extract_i_from_iuvarr, 
+            val_scale = 255.0 / N_PART_LABELS
+        )
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_or_path, filename=None, cache_dir=annotator_ckpts_path):
+        torchscript_model_path = custom_hf_download(pretrained_model_or_path, filename, cache_dir=cache_dir)
+        densepose = torch.jit.load(torchscript_model_path, map_location="cpu")
+        return cls(densepose)
+
+    def to(self, device):
+        self.dense_pose_estimation.to(device)
+        self.device = device
+        return self
+
+    def __call__(self, input_image, detect_resolution=512, output_type="pil", upscale_method="INTER_CUBIC", cmap="viridis", **kwargs):
+        input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
+        input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
+        H, W  = input_image.shape[:2]
+
+        hint_image_canvas = np.zeros([H, W], dtype=np.uint8)
+        hint_image_canvas = np.tile(hint_image_canvas[:, :, np.newaxis], [1, 1, 3])
+
+        input_image = rearrange(torch.from_numpy(input_image).to(self.device), 'h w c -> c h w')
+
+        pred_boxes, corase_segm, fine_segm, u, v = self.dense_pose_estimation(input_image)
+
+        extractor = densepose_chart_predictor_output_to_result_with_confidences
+        densepose_results = [extractor(pred_boxes[i:i+1], corase_segm[i:i+1], fine_segm[i:i+1], u[i:i+1], v[i:i+1]) for i in range(len(pred_boxes))]
+
+        hint_image = self.result_visualizer.visualize(hint_image_canvas, densepose_results)
+        hint_image = cv2.cvtColor(hint_image, cv2.COLOR_BGR2RGB)
+        hint_image = torch.from_numpy(hint_image)
+
+        if cmap=="viridis":
+            hint_image[:, :, 0][hint_image[:, :, 0] == 0] = 68
+            hint_image[:, :, 1][hint_image[:, :, 1] == 0] = 1
+            hint_image[:, :, 2][hint_image[:, :, 2] == 0] = 84
+
+        detected_map = hint_image
+
+        detected_map = detected_map.cpu().detach().numpy()
+
+        if output_type == "pil":
+            detected_map = Image.fromarray(detected_map)
+        detected_map = remove_pad(HWC3(detected_map))
+        return detected_map