initial commit
This commit is contained in:
1
.gitignore
vendored
Normal file
1
.gitignore
vendored
Normal file
@@ -0,0 +1 @@
|
||||
*.pth
|
||||
223
__init__.py
Normal file
223
__init__.py
Normal file
@@ -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": "视频剪裁"
|
||||
}
|
||||
38
model.py
Normal file
38
model.py
Normal file
@@ -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)
|
||||
112
test_single_image.py
Normal file
112
test_single_image.py
Normal file
@@ -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)
|
||||
40
utils.py
Normal file
40
utils.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user