From 5a7742e390e620eae339c4723439b68f478d1f4e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BA=B7=E5=AE=87=E4=BD=B3?= Date: Mon, 17 Feb 2025 14:48:34 +0800 Subject: [PATCH] initial commit --- .gitignore | 1 + __init__.py | 223 +++++++++++++++++++++++++++++++++++++++++++ model.py | 38 ++++++++ test_single_image.py | 112 ++++++++++++++++++++++ utils.py | 40 ++++++++ 5 files changed, 414 insertions(+) create mode 100644 .gitignore create mode 100644 __init__.py create mode 100644 model.py create mode 100644 test_single_image.py create mode 100644 utils.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b331dac --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +*.pth \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..829f0d7 --- /dev/null +++ b/__init__.py @@ -0,0 +1,223 @@ +import glob +import json +import os +import random +import shutil +import traceback +import uuid +from datetime import datetime + +from .test_single_image import test_node +import ffmpy + +video_extensions = ['webm', 'mp4', 'mkv', 'gif', 'mov'] + + +class FaceDetect: + """ + A example node + + Class methods + ------------- + INPUT_TYPES (dict): + Tell the main program input parameters of nodes. + IS_CHANGED: + optional method to control when the node is re executed. + + Attributes + ---------- + RETURN_TYPES (`tuple`): + The type of each element in the output tuple. + RETURN_NAMES (`tuple`): + Optional: The name of each output in the output tuple. + FUNCTION (`str`): + The name of the entry-point method. For example, if `FUNCTION = "execute"` then it will run Example().execute() + OUTPUT_NODE ([`bool`]): + If this node is an output node that outputs a result/image from the graph. The SaveImage node is an example. + The backend iterates on these output nodes and tries to execute all their parents if their parent graph is properly connected. + Assumed to be False if not present. + CATEGORY (`str`): + The category the node should appear in the UI. + DEPRECATED (`bool`): + Indicates whether the node is deprecated. Deprecated nodes are hidden by default in the UI, but remain + functional in existing workflows that use them. + EXPERIMENTAL (`bool`): + Indicates whether the node is experimental. Experimental nodes are marked as such in the UI and may be subject to + significant changes or removal in future versions. Use with caution in production workflows. + execute(s) -> tuple || None: + The entry point method. The name of this method must be the same as the value of property `FUNCTION`. + For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`. + """ + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + """ + Return a dictionary which contains config for all input fields. + Some types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT". + Input types "INT", "STRING" or "FLOAT" are special values for fields on the node. + The type can be a list for selection. + + Returns: `dict`: + - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required` + - Value input_fields (`dict`): Contains input fields config: + * Key field_name (`string`): Name of a entry-point method's argument + * Value field_config (`tuple`): + + First value is a string indicate the type of field or a list for selection. + + Second value is a config for type "INT", "STRING" or "FLOAT". + """ + 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" + + # OUTPUT_NODE = False + + 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: + random.seed(main_seed) + start, end = random.choice(period) + config = {"start": start, "end": end} + else: + config = {} + start = 0 + end = 0 + raise RuntimeError("未找到符合要求的视频片段") + 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 + this method can be used to force the node to execute again even when the inputs don't change. + You can make this node return a number or a string. This value will be compared to the one returned the last time the node was + executed, if it is different the node will be executed again. + This method is used in the core repo for the LoadImage node where they return the image hash as a string, if the image hash + changes between executions the LoadImage node is executed again. + """ + # @classmethod + # def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen): + # 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 VideoCut: + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "config": ("STRING", ), + "video_path": ("STRING", ), + "mod": ("INT",), + "fps": ("FLOAT",), + "period_length":("INT",{"default":10,"min":4,"max":100,"step":1,"forceInput":True}) + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("视频路径",) + + FUNCTION = "cut" + + # OUTPUT_NODE = False + + 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) + 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])]) + try: + shutil.copy(video_path, temp_fname) + except: + return ("请检查输入文件权限",) + video_path = temp_fname + # 组装输出文件名 + output_name = (".".join([*video_path.split(os.sep)[-1].split(".")[:-2], + 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(" ","")) + #调用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)), + '-segment_times', str(period_length), + '-c', 'copy', + '-map', '0', + '-avoid_negative_ts', '1' + ]} + ) + print(ff.cmd) + ff.run() + # uuid填充改回原文件名 + try: + os.remove(temp_fname) + except: + pass + try: + 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", "*")) + return (str(files),) + except: + files = glob.glob(output.replace("%03d", "*")) + traceback.print_exc() + return (str(files),) + + +# Add custom API routes, using router +from aiohttp import web +from server import PromptServer + + +@PromptServer.instance.routes.get("/hello") +async def get_hello(request): + return web.json_response("hello") + + +# A dictionary that contains all nodes you want to export with their names +# NOTE: names should be globally unique +NODE_CLASS_MAPPINGS = { + "FaceOccDetect": FaceDetect, + "VideoCutCustom": VideoCut +} + +# A dictionary that contains the friendly/humanly readable titles for the nodes +NODE_DISPLAY_NAME_MAPPINGS = { + "FaceOccDetect": "面部遮挡检测", + "VideoCutCustom": "视频剪裁" +} diff --git a/model.py b/model.py new file mode 100644 index 0000000..3477361 --- /dev/null +++ b/model.py @@ -0,0 +1,38 @@ +from torch import nn +from PIL import ImageFile + +from .utils import get_model + + +ImageFile.LOAD_TRUNCATED_IMAGES = True + + +class Model(nn.Module): + + def __init__(self, name: str, num_class: int, pretrained: bool = False, is_train: bool = True): + super(Model, self).__init__() + + self.model = get_model(name, pretrained) + + # Change the number of class + if 'resnet' in name: + in_features = self.model.fc.in_features + self.model.fc = nn.Linear(in_features, num_class) + elif 'densenet' in name: + in_features = self.model.classifier.in_features + self.model.classifier = nn.Linear(in_features, num_class) + elif "vgg" in name: + in_features = self.model.classifier[6].in_features + self.model.classifier[6] = nn.Linear(in_features, num_class) + elif "convnext" in name: + in_features = self.model.classifier[2].in_features + self.model.classifier[2] = nn.Linear(in_features, num_class) + if is_train: print(f'Model: {name}') + + def forward(self, x): + return self.model(x) + + +if __name__ == "__main__": + model = Model("convnext_large", 2, True) + print(model) \ No newline at end of file diff --git a/test_single_image.py b/test_single_image.py new file mode 100644 index 0000000..1f9efd5 --- /dev/null +++ b/test_single_image.py @@ -0,0 +1,112 @@ +import glob +import os.path +from os.path import isdir + +from PIL import Image +import torch +from torchvision import transforms +from torchvision.transforms import Resize + +from .utils import load_weight +from .model import Model + + +# CONSTANT +MEAN = [0.485, 0.456, 0.406] +STD = [0.229, 0.224, 0.225] +SIZE = [224, 224] +CLASSES = {0: "non-occluded", + 1: "occluded"} + + +def test_image(opt): + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + model = Model(opt.model, 2, False).to(device) + model = load_weight(model, opt.weight) + model.eval() + + # transform data + transform = transforms.Compose([ + transforms.Resize(SIZE), + transforms.ToTensor(), + transforms.Normalize(MEAN, STD) + ]) + + # Image + if isdir(opt.image): + imgs = glob.glob(os.path.join(opt.image,"*.png")) + for path in imgs: + img = Image.open(path).convert("RGB") + img = transform(img).to(device) + output = model(img.unsqueeze(0)) + output = torch.softmax(output, 1) + prob, pred = torch.max(output, 1) + + print("Image {} is {} - {:.2f} %".format( + path, CLASSES[pred.item()], prob.item() * 100 + )) + else: + img = Image.open(opt.image).convert("RGB") + img = transform(img).to(device) + output = model(img.unsqueeze(0)) + output = torch.softmax(output, 1) + prob, pred = torch.max(output, 1) + + print("Image {} is {} - {:.2f} %".format( + opt.image, CLASSES[pred.item()], prob.item() * 100 + )) + +def test_node(image:torch.Tensor,length=10,thres=95,model_name="convnext_tiny"): + weight_dic = { + "convnext_tiny":"best_convnext_tiny.pth", + "convnext_base":"best_convnext_base.pth" + } + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + model = Model(model_name, 2, False).to(device) + weight = os.path.join(os.path.dirname(os.path.abspath(__file__)),weight_dic[model_name]) + model = load_weight(model, weight) + model.eval() + image = image.permute(0,3,1,2) + torch_resize = Resize([224, 224]) + output = model(torch_resize(image).to(device)) + output = torch.softmax(output, 1) + prob, pred = torch.max(output, 1) + probs, preds = [round(i.item() * 100,2) for i in prob], [CLASSES[i.item()] for i in pred] + print("Image is {} - {} %".format( + preds, probs + )) + nums = [] + for idx, a,b in zip(range(len(probs)),preds,probs): + if a=="non-occluded" and b > thres: + nums.append(idx) + start = -1 + end = -1 + period = [] + for idx in range(len(nums)): + if idx == 0: + start = nums[idx] + end = nums[idx] + else: + if nums[idx] == end + 1: + end = nums[idx] + else: + if end - start + 1 >= length: + period.append([start, end]) + start = nums[idx] + end = nums[idx] + if end - start + 1 >= length: + period.append([start, end]) + return (image.permute(0,2,3,1), image.permute(0,2,3,1)[nums,:,:,:], str(preds), str(probs), str(nums), period) + + + +if __name__ == "__main__": + from argparse import ArgumentParser + + parser = ArgumentParser() + parser.add_argument("--model", type=str, help="Model name") + parser.add_argument("--weight", type=str, help="Weight path (.pth)") + parser.add_argument("--image", type=str, help="Image path") + args = parser.parse_args() + + test_image(args) diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..ca1653e --- /dev/null +++ b/utils.py @@ -0,0 +1,40 @@ +from os.path import join + +from torch import save, load +from torchvision import models + + +def save_weight(model, epoch, save_dir, file): + save({'state_dict': model.state_dict(), + 'epoch': epoch}, + join(save_dir, file)) + + +def load_weight(model, file, show=True): + checkpoints = load(file) + if show: print("Model at epoch:", checkpoints["epoch"]) + model.load_state_dict(checkpoints["state_dict"]) + return model + + +def resume_train(model, weight): + checkpoints = load(weight) + epoch = checkpoints["epoch"] + model.load_state_dict(checkpoints["state_dict"]) + return model, epoch + + + +def get_pretrained(name): + attrs = dir(models) + check = lambda x : name + "_weights" in x.lower() + # a = list(filter(check, attrs)) + + weight_class = [attr for attr in attrs if check(attr)][0] + weight = getattr(models, weight_class).IMAGENET1K_V1 + return weight + + +def get_model(name, pretrained): + model = getattr(models, name)(weights = get_pretrained(name) if pretrained else None) + return model \ No newline at end of file