142 lines
4.7 KiB
Python
142 lines
4.7 KiB
Python
import torch
|
|
from torch import Tensor
|
|
import re
|
|
from typing import Union, List, Tuple
|
|
|
|
# 定义常量
|
|
BIGMIN = -(2**53-1)
|
|
BIGMAX = 2**53-1
|
|
|
|
def convert_str_to_indexes(text: str, total_frames: int) -> List[int]:
|
|
"""将字符串转换为索引列表
|
|
|
|
支持格式:
|
|
- 单个索引: "5"
|
|
- 多个索引: "1,3,5"
|
|
- 范围: "3-10"
|
|
- 混合: "1,3-5,8,10-15"
|
|
- 负索引: "-1" (最后一帧)
|
|
- 步长: "0-10:2" (0到10每隔2帧)
|
|
"""
|
|
if not text.strip():
|
|
return []
|
|
|
|
indexes = []
|
|
parts = text.split(',')
|
|
|
|
for part in parts:
|
|
part = part.strip()
|
|
if not part:
|
|
continue
|
|
|
|
# 处理范围格式 "start-end" 或 "start-end:step"
|
|
if '-' in part and not part.startswith('-'):
|
|
# 检查是否有步长
|
|
if ':' in part:
|
|
range_part, step_str = part.split(':', 1)
|
|
try:
|
|
step = int(step_str)
|
|
except ValueError:
|
|
step = 1
|
|
else:
|
|
range_part = part
|
|
step = 1
|
|
|
|
# 解析范围
|
|
try:
|
|
start_str, end_str = range_part.split('-', 1)
|
|
start = int(start_str)
|
|
end = int(end_str)
|
|
|
|
# 处理负索引
|
|
if start < 0:
|
|
start = total_frames + start
|
|
if end < 0:
|
|
end = total_frames + end
|
|
|
|
# 确保范围有效
|
|
start = max(0, min(start, total_frames - 1))
|
|
end = max(0, min(end, total_frames - 1))
|
|
|
|
# 生成范围内的索引
|
|
if start <= end:
|
|
indexes.extend(range(start, end + 1, step))
|
|
else:
|
|
indexes.extend(range(start, end - 1, -step))
|
|
|
|
except ValueError:
|
|
# 如果解析失败，尝试作为单个索引处理
|
|
try:
|
|
idx = int(part)
|
|
if idx < 0:
|
|
idx = total_frames + idx
|
|
if 0 <= idx < total_frames:
|
|
indexes.append(idx)
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
# 处理单个索引
|
|
try:
|
|
idx = int(part)
|
|
if idx < 0:
|
|
idx = total_frames + idx
|
|
if 0 <= idx < total_frames:
|
|
indexes.append(idx)
|
|
except ValueError:
|
|
pass
|
|
|
|
# 去重并排序
|
|
indexes = sorted(list(set(indexes)))
|
|
return indexes
|
|
|
|
def select_indexes(input_obj: Tensor, indexes: List[int]) -> Tensor:
|
|
"""根据索引列表选择tensor中的数据"""
|
|
if not indexes:
|
|
return input_obj[:0] # 返回空tensor，保持维度结构
|
|
|
|
# 确保索引在有效范围内
|
|
valid_indexes = [idx for idx in indexes if 0 <= idx < input_obj.shape[0]]
|
|
|
|
if not valid_indexes:
|
|
return input_obj[:0] # 返回空tensor
|
|
|
|
return input_obj[valid_indexes]
|
|
|
|
def select_indexes_from_str(input_obj: Tensor, indexes: str,
|
|
err_if_missing: bool = True,
|
|
err_if_empty: bool = True) -> Tensor:
|
|
"""从字符串索引选择tensor数据"""
|
|
total_frames = input_obj.shape[0]
|
|
|
|
# 转换字符串为索引列表
|
|
index_list = convert_str_to_indexes(indexes, total_frames)
|
|
|
|
# 检查错误条件
|
|
if err_if_empty and not index_list:
|
|
raise ValueError(f"没有有效的索引被选择。输入: '{indexes}', 总帧数: {total_frames}")
|
|
|
|
if err_if_missing:
|
|
# 检查是否有索引超出范围
|
|
invalid_indexes = [idx for idx in index_list if idx >= total_frames or idx < 0]
|
|
if invalid_indexes:
|
|
raise ValueError(f"索引超出范围: {invalid_indexes}, 总帧数: {total_frames}")
|
|
|
|
# 选择数据
|
|
return select_indexes(input_obj, index_list)
|
|
|
|
# 描述文本
|
|
select_description = """
|
|
Image Index Selection Node
|
|
|
|
Supported Index Formats:
|
|
- Single Index: "5" (Selects the 6th image, counting from 0)
|
|
- Multiple Indexes: "1,3,5" (Selects the 2nd, 4th, and 6th images)
|
|
- Range Selection: "3-10" (Selects the 4th to 11th images)
|
|
- Range with Step: "0-10:2" (Selects the 1st, 3rd, 5th, 7th, 9th, and 11th images)
|
|
- Mixed Format: "1,3-5,8,10-15" (Selects the 1st, 3rd, 4th, 5th, 8th, 10th, 11th, 12th, 13th, 14th, and 15th images)
|
|
- Negative Index: "-1" (Selects the last image), "-3--1" (Selects the last 3 images)
|
|
|
|
Parameter Description:
|
|
- err_if_missing: Whether to report an error if the index is out of range
|
|
- err_if_empty: Whether to report an error if no images are selected
|
|
""" |