77 lines
1.9 KiB
Python
77 lines
1.9 KiB
Python
import json
|
|
|
|
from .test_single_image import test_node
|
|
|
|
|
|
class FaceDetect:
|
|
"""
|
|
人脸遮挡检测
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"main_seed": (
|
|
"INT:seed",
|
|
{"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF},
|
|
),
|
|
"model": (["convnext_tiny", "convnext_base"],),
|
|
"length": ("INT", {"default": 10, "min": 3, "max": 60, "step": 1}),
|
|
"threshold": (
|
|
"FLOAT",
|
|
{"default": 94, "min": 55, "max": 99, "step": 0.1},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"IMAGE",
|
|
"IMAGE",
|
|
"STRING",
|
|
"STRING",
|
|
"STRING",
|
|
"STRING",
|
|
"STRING",
|
|
"INT",
|
|
"INT",
|
|
)
|
|
RETURN_NAMES = (
|
|
"图像",
|
|
"选中人脸",
|
|
"分类",
|
|
"概率",
|
|
"采用帧序号",
|
|
"全部帧序列",
|
|
"剪辑配置",
|
|
"起始帧序号",
|
|
"帧数量",
|
|
)
|
|
|
|
FUNCTION = "predict"
|
|
|
|
CATEGORY = "不忘科技-自定义节点🚩/面部"
|
|
|
|
def predict(self, image, main_seed, model, length, threshold):
|
|
image, image_selected, cls, prob, nums, period = test_node(
|
|
image, length=length, thres=threshold, model_name=model
|
|
)
|
|
print("全部帧序列", period)
|
|
if len(period) > 0:
|
|
start, end = period[main_seed % len(period)]
|
|
config = {"start": start, "end": end}
|
|
else:
|
|
raise RuntimeError("未找到符合要求的视频片段")
|
|
return (
|
|
image,
|
|
image_selected,
|
|
cls,
|
|
prob,
|
|
nums,
|
|
str(period),
|
|
json.dumps(config),
|
|
start,
|
|
end - start + 1,
|
|
)
|