119 lines
3.8 KiB
Python
119 lines
3.8 KiB
Python
import os
|
|
import re
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageOps
|
|
|
|
import comfy.utils
|
|
|
|
from .utils.image_convert import pil2tensor
|
|
|
|
_CATEGORY = "KYNode/files"
|
|
|
|
class FilePathAnalyzer:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"file_path": ("STRING", {"default": "file.txt"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("parent_dir", "file_stem", "file_extension", "file_name", "full_path")
|
|
FUNCTION = "execute"
|
|
CATEGORY = _CATEGORY
|
|
DESCRIPTION = "从文件路径中提取上层目录、文件名（不含扩展名）、扩展名、文件名（含扩展名）和完整路径"
|
|
|
|
def execute(self, file_path):
|
|
# 去掉可能存在的双引号
|
|
file_path = file_path.strip('"')
|
|
path = Path(file_path)
|
|
|
|
parent_dir = str(path.parent) + "/"
|
|
file_stem = path.stem
|
|
file_extension = path.suffix
|
|
file_name = path.name
|
|
full_path = str(path.absolute())
|
|
|
|
return (parent_dir, file_stem, file_extension, file_name, full_path)
|
|
|
|
|
|
class FileSequenceAnalyzer:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"file_path": ("STRING", {"default": "image_001_v02.png"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "INT", "STRING", "INT")
|
|
RETURN_NAMES = ("sequence_pattern1", "start_number1", "sequence_pattern2", "start_number2")
|
|
FUNCTION = "execute"
|
|
CATEGORY = _CATEGORY
|
|
DESCRIPTION = """分析文件序列名称，返回两个最长的数字序列：
|
|
- sequence_pattern1: 第一个序列模板（最长序列）
|
|
- start_number1: 第一个序列起始号
|
|
- sequence_pattern2: 第二个序列模板（次长序列）
|
|
- start_number2: 第二个序列起始号
|
|
支持格式：
|
|
- image_001_v02.png -> image_###_v02.png (1) 和 image_001_v##.png (2)
|
|
- shot05_take003.jpg -> shot##_take###.jpg (5, 3)
|
|
注意：#的数量与原始数字的位数完全一致
|
|
"""
|
|
|
|
def execute(self, file_path):
|
|
path = Path(file_path.strip('"'))
|
|
file_name = path.stem # 不含扩展名的文件名
|
|
extension = path.suffix # 扩展名
|
|
|
|
# 查找文件名中的所有数字序列
|
|
# 匹配模式：文件名中的连续数字（包括前导零）
|
|
number_matches = list(re.finditer(r'(0*\d+)', file_name))
|
|
|
|
# 按序列长度排序
|
|
number_matches.sort(key=lambda x: len(x.group(1)), reverse=True)
|
|
|
|
# 准备返回值
|
|
patterns = []
|
|
numbers = []
|
|
|
|
# 处理找到的序列（最多取前两个最长的）
|
|
for match in number_matches[:2]:
|
|
number_str = match.group(1)
|
|
# 保持前导零，使用原始字符串长度
|
|
padding = len(number_str)
|
|
# 转换为数字（去掉前导零）
|
|
start_number = int(number_str)
|
|
|
|
# 构建当前序列的模式
|
|
temp_name = (
|
|
file_name[:match.start(1)] +
|
|
'#' * padding +
|
|
file_name[match.end(1):]
|
|
)
|
|
|
|
patterns.append(f"{temp_name}{extension}")
|
|
numbers.append(start_number)
|
|
|
|
# 确保始终返回两组值，不足的用空值补充
|
|
while len(patterns) < 2:
|
|
patterns.append(f"{file_name}{extension}")
|
|
numbers.append(0)
|
|
|
|
return (patterns[0], numbers[0], patterns[1], numbers[1])
|
|
|
|
|
|
FILE_CLASS_MAPPINGS = {
|
|
"KY_FilePathAnalyzer-": FilePathAnalyzer,
|
|
"KY_FileSequenceAnalyzer": FileSequenceAnalyzer,
|
|
}
|
|
|
|
FILE_NAME_MAPPINGS = {
|
|
"KY_FilePathAnalyzer-": "FilePath Analyzer",
|
|
"KY_FileSequenceAnalyzer": "File Sequence Analyzer",
|
|
}
|