From eed481f9b0bcadc7c9afe952d8d65546dd5e6480 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BA=B7=E5=AE=87=E4=BD=B3?= Date: Tue, 18 Feb 2025 17:06:49 +0800 Subject: [PATCH] =?UTF-8?q?ADD=20=E5=A2=9E=E5=8A=A0=E4=BA=BA=E8=84=B8?= =?UTF-8?q?=E6=8F=90=E5=8F=96=E8=8A=82=E7=82=B9=20ADD=20=E5=A2=9E=E5=8A=A0?= =?UTF-8?q?COS=E4=B8=8A=E4=BC=A0/=E4=B8=8B=E8=BD=BD=E8=8A=82=E7=82=B9=20PE?= =?UTF-8?q?RF=20=E5=AE=8C=E5=96=84=E9=A1=B9=E7=9B=AE=E7=BB=93=E6=9E=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 4 +- __init__.py | 192 +++++++++++++++++++++++++++++++++++++------ config.yaml | 5 ++ download/.gitkeep | 0 embedded_install.bat | 2 + install.bat | 2 + install.sh | 2 + model/.gitkeep | 0 requirements.txt | 5 +- 9 files changed, 187 insertions(+), 25 deletions(-) create mode 100644 config.yaml create mode 100644 download/.gitkeep create mode 100644 embedded_install.bat create mode 100644 install.bat create mode 100644 install.sh create mode 100644 model/.gitkeep diff --git a/.gitignore b/.gitignore index b331dac..6afb99b 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,3 @@ -*.pth \ No newline at end of file +model/*.pth +model/*.pt +!model/.gitkeep \ No newline at end of file diff --git a/__init__.py b/__init__.py index c147322..dcd3b68 100644 --- a/__init__.py +++ b/__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": "视频剪裁" } diff --git a/config.yaml b/config.yaml new file mode 100644 index 0000000..40360e3 --- /dev/null +++ b/config.yaml @@ -0,0 +1,5 @@ +secret_id: AKID63qSzJRSFyUowqa6RaVCh182jCUZS771 +secret_key: kP1ThjZDfvynQzVr9QaXHqBJOHzVowqX +region: ap-shanghai +bucket: bwkj-cos-1324682537 +subfolder: test \ No newline at end of file diff --git a/download/.gitkeep b/download/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/embedded_install.bat b/embedded_install.bat new file mode 100644 index 0000000..970df47 --- /dev/null +++ b/embedded_install.bat @@ -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 \ No newline at end of file diff --git a/install.bat b/install.bat new file mode 100644 index 0000000..69675eb --- /dev/null +++ b/install.bat @@ -0,0 +1,2 @@ + echo installing requirements + pip3.exe install -r requirements.txt \ No newline at end of file diff --git a/install.sh b/install.sh new file mode 100644 index 0000000..85f2698 --- /dev/null +++ b/install.sh @@ -0,0 +1,2 @@ + echo "installing requirements" + pip3 install -r requirements.txt \ No newline at end of file diff --git a/model/.gitkeep b/model/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/requirements.txt b/requirements.txt index 77029ba..9ec11c9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,7 @@ ffmpy torch torchvision -Pillow \ No newline at end of file +Pillow +opencv-python +ultralytics +cos-python-sdk-v5 \ No newline at end of file