ADD 增加人脸提取节点
ADD 增加COS上传/下载节点 PERF 完善项目结构
This commit is contained in:
4
.gitignore
vendored
4
.gitignore
vendored
@@ -1 +1,3 @@
|
||||
*.pth
|
||||
model/*.pth
|
||||
model/*.pt
|
||||
!model/.gitkeep
|
||||
192
__init__.py
192
__init__.py
@@ -6,6 +6,15 @@ import traceback
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
import yaml
|
||||
from ultralytics import YOLO
|
||||
from comfy import model_management
|
||||
from qcloud_cos import CosConfig, CosClientError, CosServiceError
|
||||
from qcloud_cos import CosS3Client
|
||||
|
||||
from .test_single_image import test_node
|
||||
import ffmpy
|
||||
|
||||
@@ -71,23 +80,24 @@ class FaceDetect:
|
||||
"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}),
|
||||
"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 = ("图像", "选中人脸", "分类", "概率", "采用帧序号", "全部帧序列", "剪辑配置","起始帧序号","帧数量")
|
||||
RETURN_TYPES = ("IMAGE", "IMAGE", "STRING", "STRING", "STRING", "STRING", "STRING", "INT", "INT")
|
||||
RETURN_NAMES = ("图像", "选中人脸", "分类", "概率", "采用帧序号", "全部帧序列", "剪辑配置", "起始帧序号", "帧数量")
|
||||
|
||||
FUNCTION = "predict"
|
||||
|
||||
# OUTPUT_NODE = False
|
||||
|
||||
CATEGORY = "我的自定义节点"
|
||||
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)
|
||||
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)]
|
||||
@@ -97,7 +107,7 @@ class FaceDetect:
|
||||
start = 0
|
||||
end = 0
|
||||
raise RuntimeError("未找到符合要求的视频片段")
|
||||
return (image, image_selected, cls, prob, nums, str(period), json.dumps(config), start, end-start)
|
||||
return (image, image_selected, cls, prob, nums, str(period), json.dumps(config), start, end - start)
|
||||
|
||||
"""
|
||||
The node will always be re executed if any of the inputs change but
|
||||
@@ -112,9 +122,139 @@ class FaceDetect:
|
||||
# return ""
|
||||
|
||||
|
||||
# Set the web directory, any .js file in that directory will be loaded by the frontend as a frontend extension
|
||||
# WEB_DIRECTORY = "./somejs"
|
||||
class FaceExtract():
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("图片",)
|
||||
|
||||
FUNCTION = "crop"
|
||||
|
||||
CATEGORY = "自定义节点"
|
||||
|
||||
def crop(self, image):
|
||||
device = model_management.get_torch_device()
|
||||
image_np = 255. * image.cpu().numpy()
|
||||
model = YOLO(model=os.path.join(os.path.dirname(os.path.abspath(__file__)), "model", "yolov8n-face-lindevs.pt"))
|
||||
total_images = image_np.shape[0]
|
||||
out_images = np.ndarray(shape=(total_images, 512, 512, 3))
|
||||
print("shape", image_np.shape)
|
||||
print("aaaaa")
|
||||
idx = 0
|
||||
for image_item in image_np:
|
||||
results = model.predict(
|
||||
image_item,
|
||||
imgsz=640,
|
||||
conf=0.75,
|
||||
iou=0.7,
|
||||
device=device,
|
||||
verbose=False
|
||||
)
|
||||
n = 512
|
||||
r = results[0]
|
||||
if len(r.boxes.data.cpu().numpy()) == 1:
|
||||
y1, x1, y2, x2, p, cls = r.boxes.data.cpu().numpy()[0]
|
||||
face_size = int(max(y2 - y1, x2 - x1))
|
||||
center = (x1 + x2) // 2, (y1 + y2) // 2
|
||||
x1, x2, y1, y2 = center[0] - face_size // 2, center[0] + face_size // 2, center[1] - face_size // 2, \
|
||||
center[1] + face_size // 2
|
||||
template = np.ndarray(shape=(face_size, face_size, 3))
|
||||
template.fill(20)
|
||||
for a, a1 in zip(list(range(int(x1), int(x2))), list(range(face_size))):
|
||||
for b, b1 in zip(list(range(int(y1), int(y2))), list(range(face_size))):
|
||||
if (a >= 0 and a <= r.orig_img.shape[1]) and (b >= 0 and b <= r.orig_img.shape[0]):
|
||||
template[a1][b1] = r.orig_img[a][b]
|
||||
print(int(x1), int(x2), int(y1), int(y2))
|
||||
img = cv2.resize(template, (n, n))
|
||||
out_images[idx] = img
|
||||
idx += 1
|
||||
else:
|
||||
idx += 1
|
||||
cropped_face = np.array(out_images).astype(np.float32) / 255.0
|
||||
cropped_face = torch.from_numpy(cropped_face)
|
||||
return (cropped_face,)
|
||||
|
||||
|
||||
class COSDownload:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"cos_key": ("STRING", {"multiline": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("视频存储路径",)
|
||||
FUNCTION = "download"
|
||||
CATEGORY = "自定义节点"
|
||||
|
||||
def download(self, cos_key):
|
||||
if os.sep in cos_key or "/" in cos_key or "\\" in cos_key:
|
||||
os.makedirs(os.path.join(os.path.dirname(os.path.abspath(__file__)), "download", os.path.dirname(cos_key)),
|
||||
exist_ok=True)
|
||||
for i in range(0, 10):
|
||||
try:
|
||||
with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.yaml"), encoding="utf-8",
|
||||
mode="r+") as f:
|
||||
yaml_config = yaml.load(f, Loader=yaml.FullLoader)
|
||||
config = CosConfig(Region=yaml_config["region"], SecretId=yaml_config["secret_id"],
|
||||
SecretKey=yaml_config["secret_key"])
|
||||
client = CosS3Client(config)
|
||||
response = client.download_file(
|
||||
Bucket=yaml_config["bucket"],
|
||||
Key=cos_key,
|
||||
DestFilePath=os.path.join(os.path.dirname(os.path.abspath(__file__)), "download",
|
||||
os.path.dirname(cos_key), os.path.basename(cos_key)))
|
||||
break
|
||||
except CosClientError or CosServiceError as e:
|
||||
print(f"下载失败 {e}")
|
||||
return (os.path.join(os.path.dirname(os.path.abspath(__file__)), "download", os.path.dirname(cos_key),
|
||||
os.path.basename(cos_key)),)
|
||||
|
||||
|
||||
class COSUpload:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"path": ("STRING", {"multiline": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("COS文件Key",)
|
||||
|
||||
FUNCTION = "upload"
|
||||
CATEGORY = "自定义节点"
|
||||
|
||||
def upload(self, path):
|
||||
for i in range(0, 10):
|
||||
try:
|
||||
with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.yaml"), encoding="utf-8",
|
||||
mode="r+") as f:
|
||||
yaml_config = yaml.load(f, Loader=yaml.FullLoader)
|
||||
config = CosConfig(Region=yaml_config["region"], SecretId=yaml_config["secret_id"],
|
||||
SecretKey=yaml_config["secret_key"])
|
||||
client = CosS3Client(config)
|
||||
response = client.upload_file(
|
||||
Bucket=yaml_config["bucket"],
|
||||
Key="/".join(
|
||||
[yaml_config["subfolder"], path.split("/")[-1] if "/" in path else path.split("\\")[-1]]),
|
||||
LocalFilePath=path)
|
||||
break
|
||||
except CosClientError or CosServiceError as e:
|
||||
print(e)
|
||||
return ("/".join([yaml_config["subfolder"], path.split("/")[-1] if "/" in path else path.split("\\")[-1]]),)
|
||||
|
||||
|
||||
# 有问题
|
||||
class VideoCut:
|
||||
def __init__(self):
|
||||
pass
|
||||
@@ -123,11 +263,11 @@ class VideoCut:
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"config": ("STRING", ),
|
||||
"video_path": ("STRING", ),
|
||||
"config": ("STRING",),
|
||||
"video_path": ("STRING",),
|
||||
"mod": ("INT",),
|
||||
"fps": ("FLOAT",),
|
||||
"period_length":("INT",{"default":10,"min":4,"max":100,"step":1,"forceInput":True})
|
||||
"period_length": ("INT", {"default": 10, "min": 4, "max": 100, "step": 1, "forceInput": True})
|
||||
},
|
||||
}
|
||||
|
||||
@@ -138,21 +278,21 @@ class VideoCut:
|
||||
|
||||
# OUTPUT_NODE = False
|
||||
|
||||
CATEGORY = "我的自定义节点"
|
||||
CATEGORY = "自定义节点"
|
||||
|
||||
def cut(self, config, video_path, mod, fps, period_length):
|
||||
# 原文件名
|
||||
origin_fname = ".".join(video_path.split(os.sep)[-1].split(".")[:-1])
|
||||
# 配置获取
|
||||
mul = mod/fps
|
||||
print("fps",fps)
|
||||
mul = mod / fps
|
||||
print("fps", fps)
|
||||
config = json.loads(config)
|
||||
if len(config.keys()) == 0:
|
||||
return ("无法生成符合要求的片段",)
|
||||
start, end = config["start"], config["end"]
|
||||
# 新文件名 复制改名适配ffmpeg
|
||||
uid = uuid.uuid1()
|
||||
temp_fname = os.sep.join([*video_path.split(os.sep)[:-1],"%s.%s" % (str(uid),video_path.split(".")[-1])])
|
||||
temp_fname = os.sep.join([*video_path.split(os.sep)[:-1], "%s.%s" % (str(uid), video_path.split(".")[-1])])
|
||||
try:
|
||||
shutil.copy(video_path, temp_fname)
|
||||
except:
|
||||
@@ -163,15 +303,15 @@ class VideoCut:
|
||||
video_path.split(os.sep)[-1].split(".")[-2] +
|
||||
"_output_%%03d_%s" % datetime.now().strftime('%Y%m%d_%H%M%S'),
|
||||
video_path.split(os.sep)[-1].split(".")[-1]]))
|
||||
output = (os.sep.join([*video_path.split(os.sep)[:-1],output_name])
|
||||
.replace(os.sep.join(["ComfyUI","input"]),os.sep.join(["ComfyUI","output"])).replace(" ",""))
|
||||
output = (os.sep.join([*video_path.split(os.sep)[:-1], output_name])
|
||||
.replace(os.sep.join(["ComfyUI", "input"]), os.sep.join(["ComfyUI", "output"])).replace(" ", ""))
|
||||
#调用ffmpeg
|
||||
ff = ffmpy.FFmpeg(
|
||||
inputs={video_path: ['-accurate_seek']},
|
||||
outputs={output: [
|
||||
'-f', 'segment',
|
||||
'-ss', str(round(start*mul,3)),
|
||||
'-to', str(round(end*mul,3)),
|
||||
'-ss', str(round(start * mul, 3)),
|
||||
'-to', str(round(end * mul, 3)),
|
||||
'-segment_times', str(period_length),
|
||||
'-c', 'copy',
|
||||
'-map', '0',
|
||||
@@ -186,10 +326,10 @@ class VideoCut:
|
||||
except:
|
||||
pass
|
||||
try:
|
||||
files = glob.glob(output.replace("%03d","*"))
|
||||
files = glob.glob(output.replace("%03d", "*"))
|
||||
for file in files:
|
||||
shutil.move(file,file.replace(str(uid),origin_fname))
|
||||
files = glob.glob(output.replace(str(uid),origin_fname).replace("%03d", "*"))
|
||||
shutil.move(file, file.replace(str(uid), origin_fname))
|
||||
files = glob.glob(output.replace(str(uid), origin_fname).replace("%03d", "*"))
|
||||
return (str(files),)
|
||||
except:
|
||||
files = glob.glob(output.replace("%03d", "*"))
|
||||
@@ -211,11 +351,17 @@ async def get_hello(request):
|
||||
# NOTE: names should be globally unique
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FaceOccDetect": FaceDetect,
|
||||
"FaceExtract": FaceExtract,
|
||||
"COSUpload": COSUpload,
|
||||
"COSDownload": COSDownload,
|
||||
"VideoCutCustom": VideoCut
|
||||
}
|
||||
|
||||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"FaceOccDetect": "面部遮挡检测",
|
||||
"FaceExtract": "面部提取",
|
||||
"COSUpload": "COS上传",
|
||||
"COSDownload": "COS下载",
|
||||
"VideoCutCustom": "视频剪裁"
|
||||
}
|
||||
|
||||
5
config.yaml
Normal file
5
config.yaml
Normal file
@@ -0,0 +1,5 @@
|
||||
secret_id: AKID63qSzJRSFyUowqa6RaVCh182jCUZS771
|
||||
secret_key: kP1ThjZDfvynQzVr9QaXHqBJOHzVowqX
|
||||
region: ap-shanghai
|
||||
bucket: bwkj-cos-1324682537
|
||||
subfolder: test
|
||||
0
download/.gitkeep
Normal file
0
download/.gitkeep
Normal file
2
embedded_install.bat
Normal file
2
embedded_install.bat
Normal file
@@ -0,0 +1,2 @@
|
||||
echo This will install requirements for the windows portable version of ComfyUI
|
||||
..\..\..\python_embeded\python.exe -s -m pip install -r requirements.txt
|
||||
2
install.bat
Normal file
2
install.bat
Normal file
@@ -0,0 +1,2 @@
|
||||
echo installing requirements
|
||||
pip3.exe install -r requirements.txt
|
||||
2
install.sh
Normal file
2
install.sh
Normal file
@@ -0,0 +1,2 @@
|
||||
echo "installing requirements"
|
||||
pip3 install -r requirements.txt
|
||||
0
model/.gitkeep
Normal file
0
model/.gitkeep
Normal file
@@ -1,4 +1,7 @@
|
||||
ffmpy
|
||||
torch
|
||||
torchvision
|
||||
Pillow
|
||||
Pillow
|
||||
opencv-python
|
||||
ultralytics
|
||||
cos-python-sdk-v5
|
||||
Reference in New Issue
Block a user