Merge branch 'develop'

README.md

@@ -2,7 +2,7 @@
 
 ## Installation
 
-- Install dependencies: pip install openai
+- Install dependencies: pip install -r requirements.txt
 
 - Clone the repository: git clone https://github.com/natto-maki/ComfyUI-NegiTools.git
   to your ComfyUI custom_nodes directory
@@ -76,6 +76,11 @@ Generate a noise image.
 Each component of the output image is scaled in the range of 0.0 to 1.0.
 
 
+### Generator/Depth Estimation by Marigold (experimental)
+
+Depth estimation using Marigold.
+
+
 ### utils/OpenAI Translate to English
 
 Translates text written in any language into English using GPT-4.
@@ -147,10 +152,12 @@ and fix the seed value thereafter.
 
 Outputs the properties of the image. Currently only the resolution (width and height) can be output.
 
+
 ### utils/LatentProperties
 
 Outputs the properties of the latent image. Currently only the resolution (width and height) can be output.
 
+
 ### utils/CompositeImages
 
 Composite two images with alpha.
@@ -164,3 +171,16 @@ Composite two images with alpha.
   If the alpha value is 0.0, `image_B` will be output directly; 
   if the alpha value is 1.0, the composite result will be output.
   For intermediate values, the output is the result of weighted average both images using alpha.
+
+
+### utils/OpenPoseToPointList
+
+Detects key points on the human body using OpenPose. Results are output as a JSON string.
+
+This node is used in combination with utils/PointListToMask to generate masks based on key points.
+
+
+### utils/PointListToMask
+
+Generates a mask from the coordinate list output by utils/OpenPoseToPointList.
+

__init__.py

@@ -5,6 +5,9 @@ from .negi.seed_generator import SeedGenerator
 from .negi.image_properties import ImageProperties, LatentProperties
 from .negi.composite_images import CompositeImages
 from .negi.noise_image_generator import NoiseImageGenerator
+from .negi.open_pose_to_point_list import OpenPoseToPointList
+from .negi.point_list_to_mask import PointListToMask
+from .negi.depth_estimation_by_marigold import DepthEstimationByMarigold
 
 NODE_CLASS_MAPPINGS = {
     "NegiTools_OpenAiDalle3": OpenAiDalle3,
@@ -15,6 +18,9 @@ NODE_CLASS_MAPPINGS = {
     "NegiTools_LatentProperties": LatentProperties,
     "NegiTools_CompositeImages": CompositeImages,
     "NegiTools_NoiseImageGenerator": NoiseImageGenerator,
+    "NegiTools_OpenPoseToPointList": OpenPoseToPointList,
+    "NegiTools_PointListToMask": PointListToMask,
+    "NegiTools_DepthEstimationByMarigold": DepthEstimationByMarigold,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
@@ -26,4 +32,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "NegiTools_LatentProperties": "Latent Properties 🧅",
     "NegiTools_CompositeImages": "Composite Images 🧅",
     "NegiTools_NoiseImageGenerator": "Noise Image Generator 🧅",
+    "NegiTools_OpenPoseToPointList": "OpenPose to Point List 🧅",
+    "NegiTools_PointListToMask": "Point List to Mask 🧅",
+    "NegiTools_DepthEstimationByMarigold": "Depth Estimation by Marigold (experimental) 🧅",
 }

negi/depth_estimation_by_marigold.py

@@ -0,0 +1,122 @@
+import os
+import sys
+import subprocess
+
+import numpy as np
+import torch
+import torchvision
+
+
+_dependency_dir = "dependencies"
+
+_install_script_bare = '''\
+bash script/download_weights.sh
+'''
+
+_install_script_venv = '''\
+source venv/marigold/bin/activate
+pip install -r requirements.txt
+bash script/download_weights.sh
+'''
+
+_infer_script_bare = '''\
+%(interpreter)s run.py --n_infer %(infer_passes)d --denoise_steps %(denoise_steps)d --seed %(seed)d --input_rgb_dir "%(input_dir_name)s" --output_dir "%(output_dir_name)s"
+'''
+
+_infer_script_venv = '''\
+source venv/marigold/bin/activate
+python run.py --n_infer %(infer_passes)d --denoise_steps %(denoise_steps)d --seed %(seed)d --input_rgb_dir "%(input_dir_name)s" --output_dir "%(output_dir_name)s"
+'''
+
+
+class DepthEstimationByMarigold:
+    def __check_environment(self, enable_venv=False):
+        if not os.path.isdir(os.path.join(self.dep_dir, "Marigold")):
+            r0 = subprocess.run(["git", "clone", "https://github.com/prs-eth/Marigold.git"], cwd=self.dep_dir)
+            if r0.returncode != 0:
+                subprocess.run(["rm", "-rf", "Marigold"], cwd=self.dep_dir)
+                raise RuntimeError("Marigold repository not found or connection error")
+
+        if not enable_venv and not os.path.isfile(os.path.join(self.rep_dir, "installed_bare")):
+            with open(os.path.join(self.rep_dir, "install.sh"), "wt") as f:
+                f.write(_install_script_bare)
+            subprocess.run(["bash", "install.sh"], cwd=self.rep_dir)
+            with open(os.path.join(self.rep_dir, "installed_bare"), "wt") as f:
+                f.write("installed")
+
+        if enable_venv and not os.path.isfile(os.path.join(self.rep_dir, "installed_venv")):
+            # Make sure that venv has been created correctly.
+            # Because if you ignore the error, the ComfyUI runtime environment package will be incorrectly overwritten.
+            subprocess.run([sys.executable, "-m", "venv", "venv/marigold"], cwd=self.rep_dir)
+            if not os.path.isfile(os.path.join(self.rep_dir, "venv", "marigold", "bin", "activate")):
+                raise RuntimeError("Failed to setup venv for Marigold")
+
+            with open(os.path.join(self.rep_dir, "install.sh"), "wt") as f:
+                f.write(_install_script_venv)
+            # TODO pick errors
+            subprocess.run(["bash", "install.sh"], cwd=self.rep_dir)
+            with open(os.path.join(self.rep_dir, "installed_venv"), "wt") as f:
+                f.write("installed")
+
+
+    def __init__(self):
+        self.dep_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), _dependency_dir)
+        os.makedirs(self.dep_dir, exist_ok=True)
+        self.rep_dir = os.path.join(self.dep_dir, "Marigold")
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "infer_passes": ("INT", {"default": 10, "min": 1, "max": 40, "step": 1, "display": "number"}),
+                "denoise_steps": ("INT", {"default": 10, "min": 1, "max": 40, "step": 1, "display": "number"}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffff}),
+                "runtime": ([
+                   "bare (recommended)",
+                   "venv (if \"bare\" doesn't work)",
+                ],),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("DEPTH_IMAGE",)
+    FUNCTION = "doit"
+    OUTPUT_NODE = False
+    CATEGORY = "Generator"
+
+    def doit(self, image, infer_passes, denoise_steps, seed, runtime):
+        use_venv = runtime.startswith("venv")
+        self.__check_environment(use_venv)
+
+        work_dir = os.path.join(self.rep_dir, "work")
+        os.makedirs(work_dir, exist_ok=True)
+
+        input_dir = os.path.join(work_dir, "input")
+        output_dir = os.path.join(work_dir, "output")
+        subprocess.run(["rm", "-rf", "input"], cwd=work_dir)
+        os.makedirs(input_dir, exist_ok=True)
+
+        im0 = torchvision.transforms.functional.to_pil_image(torch.permute(image[0], (2, 0, 1)))
+        im0.save(os.path.join(input_dir, "image.png"))
+
+        with open(os.path.join(work_dir, "infer.sh"), "wt") as f:
+            f.write((_infer_script_venv if use_venv else _infer_script_bare) % {
+                "interpreter": os.path.abspath(sys.executable),
+                "infer_passes": infer_passes,
+                "denoise_steps": denoise_steps,
+                "seed": seed,
+                "input_dir_name": os.path.abspath(input_dir),
+                "output_dir_name": os.path.abspath(output_dir)
+            })
+
+        subprocess.run(["bash", os.path.join("work", "infer.sh")], cwd=self.rep_dir)
+
+        im1 = np.load(os.path.join(output_dir, "depth_npy", "image_pred.npy")).astype(np.float32)
+        im1_min = np.min(im1)
+        im1_max = np.max(im1)
+        if im1_min == im1_max:
+            im1_max = im1_min + 1.0
+        im1 = (im1 - im1_min) * (1.0 / (im1_max - im1_min))
+
+        return (torch.from_numpy(np.expand_dims(np.stack([im1, im1, im1], axis=-1), axis=0)),)

negi/open_pose_to_point_list.py

@@ -0,0 +1,86 @@
+import json
+
+import numpy as np
+from controlnet_aux import OpenposeDetector
+from controlnet_aux.util import HWC3, resize_image
+
+
+_names = [
+    "Nose", "Neck",
+    "RShoulder", "RElbow", "RWrist",
+    "LShoulder", "LElbow", "LWrist",
+    "RHip", "RKnee", "RAnkle",
+    "LHip", "LKnee", "LAnkle",
+    "REye", "LEye", "REar", "LEar"
+]
+
+_name_to_index = {name: i for i, name in enumerate(_names)}
+
+
+class OpenPoseToPointList:
+    def __init__(self):
+        self.open_pose = OpenposeDetector.from_pretrained("lllyasviel/Annotators")
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "detect_resolution": ("INT", {"default": 512, "min": 64, "max": 2048, "step": 64, "display": "slider"}),
+                "method": ([
+                   "face",
+                   "hand",
+                   "all",
+                ],),
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("POINT_LIST",)
+    FUNCTION = "doit"
+    OUTPUT_NODE = False
+    CATEGORY = "utils"
+
+    def doit(self, image, detect_resolution, method):
+        input_image = (np.fmax(0.0, np.fmin(1.0, image.to('cpu').detach().numpy()[0])) * 255.0).astype(np.uint8)
+        input_image = HWC3(input_image)
+        input_image = resize_image(input_image, detect_resolution)
+
+        poses = self.open_pose.detect_poses(input_image, include_hand=False, include_face=False)
+
+        if method == "face":
+            ret = []
+            for pose in poses:
+                x = 0.0
+                y = 0.0
+                n = 0
+                for name in ["Nose", "REye", "LEye", "REar", "LEar"]:
+                    key_point = pose.body.keypoints[_name_to_index[name]]
+                    if key_point is not None:
+                        x += key_point.x
+                        y += key_point.y
+                        n += 1
+                if n != 0:
+                    ret.append({"x": x / n, "y": y / n})
+
+        elif method == "hand":
+            ret = []
+            for pose in poses:
+                for name in ["RWrist", "LWrist"]:
+                    key_point = pose.body.keypoints[_name_to_index[name]]
+                    if key_point is not None:
+                        ret.append({"x": key_point.x, "y": key_point.y})
+
+        elif method == "all":
+            ret = []
+            for pose in poses:
+                points = {}
+                for i, key_point in enumerate(pose.body.keypoints):
+                    if key_point is not None:
+                        points[_names[i]] = {"x": key_point.x, "y": key_point.y, "score": key_point.score}
+                ret.append(points)
+
+        else:
+            raise ValueError()
+
+        return (json.dumps(ret, indent=2),)

negi/point_list_to_mask.py

@@ -0,0 +1,41 @@
+import json
+
+import numpy as np
+
+
+class PointListToMask:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "point_list": ("STRING", {"multiline": False, "default": ""}),
+                "width": ("INT", {"default": 512, "min": 0, "max": 4096, "step": 64, "display": "number"}),
+                "height": ("INT", {"default": 512, "min": 0, "max": 4096, "step": 64, "display": "number"}),
+                "radius": ("INT", {"default": 50, "min": 1, "max": 2048, "step": 1, "display": "number"}),
+            }
+        }
+
+    RETURN_TYPES = ("MASK",)
+    FUNCTION = "doit"
+    OUTPUT_NODE = False
+    CATEGORY = "utils"
+
+    def doit(self, point_list, width, height, radius):
+        point_list = json.loads(point_list)
+
+        ret = []
+        px = (np.reshape(np.arange(width, dtype=np.float32), (1, -1))
+              * np.ones((height, 1), dtype=np.float32))
+        py = (np.reshape(np.arange(height, dtype=np.float32), (-1, 1))
+              * np.ones((1, width), dtype=np.float32))
+        for point in point_list:
+            d2 = np.power(px - point["x"] * width, 2.0) + np.power(py - point["y"] * height, 2.0)
+            ret.append(np.reshape(d2 <= radius * radius, (height, width)).astype(np.float32))
+
+        if len(ret) == 0:
+            return (np.zeros((1, height, width), dtype=np.float32),)
+
+        return (np.fmin(1.0, np.sum(np.stack(ret), axis=0, keepdims=True)),)

requirements.txt

@@ -0,0 +1,2 @@
+openai >= 1.3.0
+controlnet-aux >= 0.0.7