ADD 增加人脸提取节点

ADD 增加COS上传/下载节点
PERF 完善项目结构

.gitignore

@@ -1 +1,3 @@
-*.pth
+model/*.pth
+model/*.pt
+!model/.gitkeep

__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
 
@@ -84,10 +93,11 @@ class FaceDetect:
 
     # 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)]
@@ -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
+class FaceExtract():
-# WEB_DIRECTORY = "./somejs"
+    @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
@@ -138,7 +278,7 @@ class VideoCut:
 
     # OUTPUT_NODE = False
 
-    CATEGORY = "我的自定义节点"
+    CATEGORY = "自定义节点"
 
     def cut(self, config, video_path, mod, fps, period_length):
         # 原文件名
@@ -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": "视频剪裁"
 }

config.yaml

@@ -0,0 +1,5 @@
+secret_id: AKID63qSzJRSFyUowqa6RaVCh182jCUZS771
+secret_key: kP1ThjZDfvynQzVr9QaXHqBJOHzVowqX
+region: ap-shanghai
+bucket: bwkj-cos-1324682537
+subfolder: test

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

install.bat

@@ -0,0 +1,2 @@
+ echo installing requirements
+ pip3.exe install -r requirements.txt

install.sh

@@ -0,0 +1,2 @@
+ echo "installing requirements"
+ pip3 install -r requirements.txt

requirements.txt

@@ -2,3 +2,6 @@ ffmpy
 torch
 torchvision
 Pillow
+opencv-python
+ultralytics
+cos-python-sdk-v5