Patrice Ferlet
069eafb8d8
Making hot fixes is a bad idea, I badly merged and fix a conflict without correctly testing.
73 lines
1.9 KiB
Python
73 lines
1.9 KiB
Python
from typing import Dict, Tuple
|
|
|
|
import onnxruntime as ort
|
|
import torch
|
|
|
|
from .detector.human_parts import get_mask, labels
|
|
from .utils import model_path
|
|
|
|
|
|
class HumanParts:
|
|
"""
|
|
This node is used to get a mask of the human parts in the image.
|
|
|
|
The model used is DeepLabV3+ with a ResNet50 backbone trained
|
|
by Keras-io, converted to ONNX format.
|
|
|
|
"""
|
|
|
|
RETURN_TYPES = ("MASK",)
|
|
RETURN_NAMES = ("mask",)
|
|
FUNCTION = "get_mask"
|
|
CATEGORY = "Metal3d"
|
|
OUTPU_NODE = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
def _bool_widget(
|
|
is_enabled=False, tooltip: str | None = None
|
|
) -> Tuple[str, dict]:
|
|
"""Helper function to create a boolean widget"""
|
|
return (
|
|
"BOOLEAN",
|
|
{
|
|
"default": is_enabled,
|
|
"label_on": "Enabled",
|
|
"label_off": "Disabled",
|
|
"tooltip": tooltip,
|
|
},
|
|
)
|
|
|
|
# automate the creation of the inputs using the known labels
|
|
entries: Dict[str, tuple] = {
|
|
segment[0]: _bool_widget(False, f"{segment[1]}")
|
|
for segment in labels.values()
|
|
if segment[0] != ""
|
|
}
|
|
|
|
inputs = {
|
|
"required": {
|
|
"image": (
|
|
"IMAGE",
|
|
{
|
|
"label": "Image",
|
|
"tooltip": "The image in which to detect human parts",
|
|
},
|
|
)
|
|
},
|
|
"optional": {},
|
|
}
|
|
inputs["required"].update(entries)
|
|
|
|
return inputs
|
|
|
|
def get_mask(self, image: torch.Tensor, **kwargs) -> Tuple[torch.Tensor]:
|
|
"""
|
|
Return a Tensor with the mask of the human parts in the image.
|
|
"""
|
|
|
|
model = ort.InferenceSession(model_path)
|
|
ret_tensor, _ = get_mask(image, model=model, rotation=0, **kwargs)
|
|
|
|
return (ret_tensor,)
|