comi-logic/ComfyUI-CustomNode
8
0
Fork 0
Code Pull Requests Releases Activity

fix: 判断文件是否存在

Browse Source
This commit is contained in:
ymm
2025-02-19 21:50:18 +08:00
parent fc42888f5a
commit f00e51b152
1 changed files with 212 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

285
__init__.py
View File

@@ -18,7 +18,7 @@ from qcloud_cos import CosS3Client
from .test_single_image import test_node
import ffmpy
video_extensions = ['webm', 'mp4', 'mkv', 'gif', 'mov']
video_extensions = ["webm", "mp4", "mkv", "gif", "mov"]
class FaceDetect:
@@ -31,30 +31,67 @@ class FaceDetect:
return {
"required": {
"image": ("IMAGE",),
"main_seed": ("INT:seed", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
"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})
"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"
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)]
config = {"start": start, "end": end}
else:
raise RuntimeError("未找到符合要求的视频片段")
return (image, image_selected, cls, prob, nums, str(period), json.dumps(config), start, end - start+1)
return (
image,
image_selected,
cls,
prob,
nums,
str(period),
json.dumps(config),
start,
end - start + 1,
)
class FaceExtract:
@@ -77,8 +114,14 @@ class FaceExtract:
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"))
image_np = 255.0 * 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)
@@ -86,12 +129,7 @@ class FaceExtract:
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
image_item, imgsz=640, conf=0.75, iou=0.7, device=device, verbose=False
)
n = 512
r = results[0]
@@ -99,13 +137,21 @@ class FaceExtract:
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
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]):
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))
@@ -137,26 +183,51 @@ class COSDownload:
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)
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:
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"])
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)))
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)),)
return (
os.path.join(
os.path.dirname(os.path.abspath(__file__)),
"download",
os.path.dirname(cos_key),
os.path.basename(cos_key),
),
)
class COSUpload:
@@ -181,21 +252,45 @@ class COSUpload:
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:
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"])
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)
[
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]]),)
return (
"/".join(
[
yaml_config["subfolder"],
path.split("/")[-1] if "/" in path else path.split("\\")[-1],
]
),
)
class VideoCut:
@@ -209,7 +304,16 @@ class VideoCut:
"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,
},
),
},
}
@@ -234,31 +338,54 @@ class VideoCut:
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:
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
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'
]}
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()
@@ -271,7 +398,9 @@ class VideoCut:
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", "*"))
files = glob.glob(
output.replace(str(uid), origin_fname).replace("%03d", "*")
)
return (str(files),)
except:
files = glob.glob(output.replace("%03d", "*"))
@@ -289,8 +418,6 @@ async def get_hello(request):
return web.json_response("hello")
# 腾讯云 VOD
from tencentcloud.common.profile.client_profile import ClientProfile
@@ -301,13 +428,14 @@ import requests
from pathlib import Path
import tempfile
class VodToLocalNode:
def __init__(self):
self.secret_id = "AKIDsrihIyjZOBsjimt8TsN8yvv1AMh5dB44"
self.secret_key = "CPZcxdk6W39Jd4cGY95wvupoyMd0YFqW"
self.vod_client = self.init_vod_client()
def init_vod_client(self):
"""初始化VOD客户端"""
try:
@@ -336,22 +464,22 @@ class VodToLocalNode:
# 调用下载逻辑
local_path = self.download_vod(file_id, sub_app_id)
return (local_path,)
def _get_download_url(self, file_id, sub_app_id):
"""获取媒体文件下载地址"""
try:
req = models.DescribeMediaInfosRequest()
req.FileIds = [file_id]
req.SubAppId = int(sub_app_id)
resp = self.vod_client.DescribeMediaInfos(req)
if not resp.MediaInfoSet:
raise ValueError("File not found")
media_info = resp.MediaInfoSet[0]
if not media_info.BasicInfo.MediaUrl:
raise ValueError("No download URL available")
return media_info.BasicInfo.MediaUrl
except Exception as e:
raise RuntimeError(f"Tencent API error: {e}")
@@ -359,7 +487,9 @@ class VodToLocalNode:
def create_directory(self, path):
p = Path(path)
if not p.exists():
p.mkdir(parents=True, exist_ok=True) # parents=True会自动创建所有必需的父目录，exist_ok=True表示如果目录已存在则不会引发异常
p.mkdir(
parents=True, exist_ok=True
) # parents=True会自动创建所有必需的父目录，exist_ok=True表示如果目录已存在则不会引发异常
print(f"目录已创建: {path}")
else:
print(f"目录已存在: {path}")
@@ -371,14 +501,23 @@ class VodToLocalNode:
"""
media_url = self._get_download_url(file_id=file_id, sub_app_id=sub_app_id)
# 生成一个临时目录路径名并创建该目录
self.create_directory(os.path.join(os.path.dirname(os.path.abspath(__file__)), "download", f"{sub_app_id}"))
output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "download", f"{sub_app_id}", f"{file_id}.mp4" )
self.create_directory(
os.path.join(
os.path.dirname(os.path.abspath(__file__)), "download", f"{sub_app_id}"
)
)
output_dir = os.path.join(
os.path.dirname(os.path.abspath(__file__)),
"download",
f"{sub_app_id}",
f"{file_id}.mp4",
)
# 判断文件是否存在
if os.path.exists(output_dir):
return output_dir
return self._download_file(url=media_url, save_path=output_dir, timeout=60 * 10)
def _download_file(self, url: str, save_path: str, timeout: int=30):
def _download_file(self, url: str, save_path: str, timeout: int = 30):
"""下载文件到本地"""
try:
with requests.get(url, stream=True, timeout=timeout) as response:
@@ -391,6 +530,7 @@ class VodToLocalNode:
except Exception as e:
raise RuntimeError(f"Download error: {e}")
# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
NODE_CLASS_MAPPINGS = {
@@ -399,7 +539,7 @@ NODE_CLASS_MAPPINGS = {
"COSUpload": COSUpload,
"COSDownload": COSDownload,
"VideoCutCustom": VideoCut,
"VodToLocal": VodToLocalNode
"VodToLocal": VodToLocalNode,
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -409,6 +549,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"COSUpload": "COS上传",
"COSDownload": "COS下载",
"VideoCutCustom": "视频剪裁",
"VodToLocal": "腾讯云VOD下载"
"VodToLocal": "腾讯云VOD下载",
}