774 lines
29 KiB
Python
774 lines
29 KiB
Python
import os
|
|
import random
|
|
import math
|
|
from PIL import Image, ImageDraw,ImageOps
|
|
import torch
|
|
from typing import Union
|
|
import numpy as np
|
|
from pathlib import Path
|
|
# from torchvision import transforms
|
|
|
|
# from .tensor_pillow_loading import pillow_image_from_path
|
|
# from .tensor_pillow_conversion import pillow_image_from_tensor
|
|
|
|
def pillow_image_from_path(image_path: str) -> Image.Image:
|
|
"""
|
|
从指定路径加载图像为 Pillow 图像。
|
|
:param image_path: 图像文件的路径
|
|
:return: 加载后的 Pillow 图像
|
|
"""
|
|
image = Image.open(image_path)
|
|
# print(f"图像 {image_path} 已从文件加载为 Pillow 图像。 | Image {image_path} has been loaded from file as a Pillow image.")
|
|
return image
|
|
|
|
def pillow_image_from_tensor(tensor: torch.Tensor):
|
|
# 先将张量移到 CPU 并转换为 numpy 数组
|
|
numpy_array = tensor.cpu().numpy()
|
|
|
|
# 确保数组的维度和数据类型正确
|
|
if numpy_array.ndim == 4:
|
|
# 如果是批量数据，取第一个图像
|
|
numpy_array = numpy_array[0]
|
|
|
|
# 调整维度顺序，从 (C, H, W) 到 (H, W, C)
|
|
if numpy_array.shape[0] in [3, 4]:
|
|
numpy_array = np.transpose(numpy_array, (1, 2, 0))
|
|
|
|
# 将数据类型转换为 uint8
|
|
if numpy_array.dtype != np.uint8:
|
|
# 如果数据范围是 [0, 1]，将其转换为 [0, 255]
|
|
if numpy_array.max() <= 1.0:
|
|
numpy_array = (numpy_array * 255).astype(np.uint8)
|
|
else:
|
|
numpy_array = numpy_array.astype(np.uint8)
|
|
|
|
# 创建 Pillow 图像
|
|
image = Image.fromarray(numpy_array)
|
|
return image
|
|
|
|
|
|
|
|
def pillow_image_get_shape(image:Image.Image)-> tuple[int,int]:
|
|
width, height = image.size
|
|
return tuple(width, height)
|
|
|
|
def pillow_image_to_gray(image:Image.Image)->Image.Image:
|
|
return image.convert('L')
|
|
|
|
def pillow_image_to_pillow_mask(image:Image.Image, channel:str):
|
|
"""
|
|
将 Pillow 图像转换为 Pillow 掩码。
|
|
|
|
:param pillow_image: 输入的 Pillow 图像对象
|
|
:param channel: 选择的通道，字符串类型，可选值为 'red', 'green', 'blue', 'alpha'
|
|
:return: 转换后的 Pillow 掩码图像对象
|
|
"""
|
|
channels = ["red", "green", "blue", "alpha"]
|
|
if channel not in channels:
|
|
raise ValueError(f"无效的通道名称 '{channel}'，请选择 'red', 'green', 'blue' 或 'alpha'。")
|
|
|
|
# 将 Pillow 图像转换为 numpy 数组
|
|
|
|
image_array = np.array(image)
|
|
|
|
# 根据通道名称确定索引
|
|
channel_index = channels.index(channel)
|
|
if len(image_array.shape) == 3: # 彩色图像
|
|
mask_array = image_array[:, :, channel_index]
|
|
elif len(image_array.shape) == 2: # 灰度图像
|
|
if channel != 'alpha':
|
|
mask_array = image_array
|
|
else:
|
|
raise ValueError("灰度图像没有 alpha 通道。")
|
|
else:
|
|
raise ValueError("输入的图像数组维度不符合要求。")
|
|
|
|
# 将 numpy 数组转换回 Pillow 图像
|
|
mask = Image.fromarray(mask_array)
|
|
return mask
|
|
|
|
|
|
def pillow_image_from_input(input: Union[str, torch.Tensor,Image.Image], is_mask: bool = False) -> Image.Image:
|
|
"""
|
|
从图像路径或图像张量加载Pillow图像或遮罩。
|
|
|
|
Args:
|
|
input (Union[str, torch.Tensor]): 输入可以是图像/遮罩的文件路径（字符串），
|
|
或者是一个图像/遮罩张量（torch.Tensor）。
|
|
is_mask (bool): 如果为True，则表示加载的是遮罩；否则为图像。默认值为False。
|
|
|
|
Returns:
|
|
Image.Image: 加载后的Pillow图像或遮罩对象。
|
|
|
|
Raises:
|
|
ValueError: 如果输入既不是字符串也不是torch.Tensor类型。
|
|
FileNotFoundError: 如果输入是字符串但文件不存在。
|
|
"""
|
|
|
|
if isinstance(input, str):
|
|
image = pillow_image_from_path(input)
|
|
if is_mask:
|
|
image = pillow_image_to_gray(image)
|
|
return image
|
|
|
|
if isinstance(input, torch.Tensor):
|
|
image =pillow_image_from_tensor(input)
|
|
if is_mask:
|
|
image = pillow_image_to_gray(image)
|
|
return image
|
|
|
|
if isinstance(input, Image.Image):
|
|
image =input
|
|
if is_mask:
|
|
image = pillow_image_to_gray(image)
|
|
return image
|
|
|
|
|
|
|
|
def create_shape_mask(shape: str, width: int, height: int) -> Image.Image:
|
|
"""
|
|
创建指定形状的遮罩
|
|
|
|
Args:
|
|
shape (str): 形状名称，如 "square", "triangle", "circle", "heart", "star", "diamond"
|
|
width (int): 遮罩的宽度
|
|
height (int): 遮罩的高度
|
|
|
|
Returns:
|
|
Image.Image: 创建好的遮罩图像
|
|
"""
|
|
mask = Image.new('L', (width, height), 0)
|
|
draw = ImageDraw.Draw(mask)
|
|
if shape == "square":
|
|
draw.rectangle((0, 0, width, height), fill=255)
|
|
elif shape == "triangle":
|
|
draw.polygon([(0, height), (width // 2, 0), (width, height)], fill=255)
|
|
elif shape == "circle":
|
|
draw.ellipse((0, 0, width, height), fill=255)
|
|
elif shape == "heart":
|
|
points = []
|
|
for t in range(0, 360, 10):
|
|
rad = math.radians(t)
|
|
x_coord = 16 * math.sin(rad) ** 3
|
|
y_coord = -(13 * math.cos(rad) - 5 * math.cos(2 * rad) - 2 * math.cos(3 * rad) - math.cos(4 * rad))
|
|
x_scaled = int(x_coord * width / 32 + width / 2)
|
|
y_scaled = int(y_coord * height / 32 + height / 2)
|
|
points.append((x_scaled, y_scaled))
|
|
draw.polygon(points, fill=255)
|
|
elif shape == "star":
|
|
outer_radius = min(width, height) // 2
|
|
inner_radius = outer_radius // 2
|
|
center_x = width // 2
|
|
center_y = height // 2
|
|
points = []
|
|
for i in range(5):
|
|
angle1 = 2 * math.pi * i / 5
|
|
angle2 = 2 * math.pi * (i + 0.5) / 5
|
|
x1 = center_x + outer_radius * math.cos(angle1)
|
|
y1 = center_y + outer_radius * math.sin(angle1)
|
|
x2 = center_x + inner_radius * math.cos(angle2)
|
|
y2 = center_y + inner_radius * math.sin(angle2)
|
|
points.extend([(x1, y1), (x2, y2)])
|
|
draw.polygon(points, fill=255)
|
|
elif shape == "diamond":
|
|
draw.polygon([(width // 2, 0), (width, height // 2), (width // 2, height), (0, height // 2)], fill=255)
|
|
return mask
|
|
|
|
def get_shard_dimensions(x: int, y: int, shard_size: int, width: int, height: int) -> tuple:
|
|
"""
|
|
计算单个碎片的实际尺寸
|
|
|
|
Args:
|
|
x (int): 碎片的起始 x 坐标
|
|
y (int): 碎片的起始 y 坐标
|
|
shard_size (int): 碎片的大小
|
|
width (int): 原始图片的宽度
|
|
height (int): 原始图片的高度
|
|
|
|
Returns:
|
|
tuple: 碎片的实际宽度和高度
|
|
"""
|
|
actual_width = min(shard_size, width - x)
|
|
actual_height = min(shard_size, height - y)
|
|
return actual_width, actual_height
|
|
|
|
def extract_shard_image(original_img: Image.Image, x: int, y: int, actual_width: int, actual_height: int) -> Image.Image:
|
|
"""
|
|
从原始图片中提取单个碎片图像
|
|
|
|
Args:
|
|
original_img (Image.Image): 原始图片
|
|
x (int): 碎片的起始 x 坐标
|
|
y (int): 碎片的起始 y 坐标
|
|
actual_width (int): 碎片的实际宽度
|
|
actual_height (int): 碎片的实际高度
|
|
|
|
Returns:
|
|
Image.Image: 提取的碎片图像
|
|
"""
|
|
return original_img.crop((x, y, x + actual_width, y + actual_height))
|
|
|
|
def apply_shard_mask(shard_img: Image.Image, shard_shape: str, shard_mask_image: Union[str, torch.Tensor], actual_width: int, actual_height: int) -> Image.Image:
|
|
"""
|
|
为碎片图像应用遮罩形状和遮罩图像
|
|
|
|
Args:
|
|
shard_img (Image.Image): 碎片图像
|
|
shard_shape (str): 碎片的形状
|
|
shard_mask_image (Union[str, torch.Tensor]): 碎片的遮罩图像
|
|
actual_width (int): 碎片的实际宽度
|
|
actual_height (int): 碎片的实际高度
|
|
|
|
Returns:
|
|
Image.Image: 应用遮罩后的碎片图像
|
|
"""
|
|
if shard_shape:
|
|
mask = create_shape_mask(shard_shape, actual_width, actual_height)
|
|
shard_img.putalpha(mask)
|
|
if shard_mask_image is not None:
|
|
# Boolean value of Tensor with more than one value is ambiguous
|
|
mask = pillow_image_from_input(shard_mask_image, is_mask=True)
|
|
mask = mask.resize((actual_width, actual_height))
|
|
shard_img.putalpha(mask)
|
|
return shard_img
|
|
|
|
def generate_shard_movement_params(speed: int, rotation: int, frame_count: int, fade_start: int) -> dict:
|
|
"""
|
|
生成单个碎片的运动参数
|
|
|
|
Args:
|
|
speed (int): 碎片的运动速度
|
|
rotation (int): 碎片的旋转速度
|
|
frame_count (int): 总帧数
|
|
fade_start (int): 碎片开始消失的帧数
|
|
|
|
Returns:
|
|
dict: 碎片的运动参数
|
|
"""
|
|
angle = random.uniform(0, 2 * math.pi)
|
|
shard_speed = random.uniform(0.5, 1.5) * speed
|
|
rot_speed = random.uniform(-1, 1) * rotation
|
|
fade = random.randint(frame_count // 2, frame_count - 5) if max(0, fade_start) < 1 else random.randint(max(0, fade_start - 5), min(frame_count - 1, fade_start + 5))
|
|
return {
|
|
'angle': angle,
|
|
'speed': shard_speed,
|
|
'rotation': 0,
|
|
'rotation_speed': rot_speed,
|
|
'fade_start': fade,
|
|
'opacity': 1.0
|
|
}
|
|
|
|
def create_shard_info(shard_img: Image.Image, x: int, y: int, speed: int, rotation: int, frame_count: int, fade_start: int) -> dict:
|
|
"""
|
|
创建单个碎片的完整信息
|
|
|
|
Args:
|
|
shard_img (Image.Image): 碎片图像
|
|
x (int): 碎片的起始 x 坐标
|
|
y (int): 碎片的起始 y 坐标
|
|
speed (int): 碎片的运动速度
|
|
rotation (int): 碎片的旋转速度
|
|
frame_count (int): 总帧数
|
|
fade_start (int): 碎片开始消失的帧数
|
|
|
|
Returns:
|
|
dict: 碎片的完整信息
|
|
"""
|
|
movement_params = generate_shard_movement_params(speed, rotation, frame_count, fade_start)
|
|
return {
|
|
'img': shard_img,
|
|
'x': x,
|
|
'y': y,
|
|
'width': shard_img.width,
|
|
'height': shard_img.height,
|
|
**movement_params
|
|
}
|
|
|
|
def update_shard_position_and_rotation(shard: dict, frame: int):
|
|
"""
|
|
更新单个碎片的位置和旋转
|
|
|
|
Args:
|
|
shard (dict): 碎片的完整信息
|
|
frame (int): 当前帧数
|
|
"""
|
|
if frame > 0:
|
|
shard['x'] += math.cos(shard['angle']) * shard['speed']
|
|
shard['y'] += math.sin(shard['angle']) * shard['speed']
|
|
shard['rotation'] += shard['rotation_speed']
|
|
|
|
def update_shard_opacity(shard: dict, frame: int):
|
|
"""
|
|
更新单个碎片的透明度
|
|
|
|
Args:
|
|
shard (dict): 碎片的完整信息
|
|
frame (int): 当前帧数
|
|
"""
|
|
if frame >= shard['fade_start']:
|
|
shard['opacity'] = max(0, shard['opacity'] - 0.1)
|
|
|
|
def rotate_and_apply_opacity(shard: dict) -> Image.Image:
|
|
"""
|
|
旋转碎片图像并应用透明度
|
|
|
|
Args:
|
|
shard (dict): 碎片的完整信息
|
|
|
|
Returns:
|
|
Image.Image: 旋转并应用透明度后的碎片图像
|
|
"""
|
|
rotated_shard = shard['img'].rotate(shard['rotation'], expand=True)
|
|
if shard['opacity'] < 1.0:
|
|
alpha = rotated_shard.split()[3]
|
|
alpha = alpha.point(lambda p: p * shard['opacity'])
|
|
rotated_shard.putalpha(alpha)
|
|
return rotated_shard
|
|
|
|
def text_padd(text: str, action: str, length: int, filled: str = ' ') -> str:
|
|
"""
|
|
对字符串进行填充操作
|
|
|
|
Args:
|
|
text (str): 要填充的字符串
|
|
action (str): 填充方向，'padd-start' 或 'padd-end'
|
|
length (int): 填充后的总长度
|
|
filled (str): 填充字符，默认为空格
|
|
|
|
Returns:
|
|
str: 填充后的字符串
|
|
"""
|
|
res = text
|
|
if 'padd-start' in action:
|
|
res = text.rjust(length, filled)
|
|
elif 'padd-end' in action:
|
|
res = text.ljust(length, filled)
|
|
return res
|
|
|
|
# 编写一个函数，将 false,true转为 bool。忽略大小写。支持传入 字符或者布尔。添加函数提示和注释。
|
|
# 支持 enable,disable,1,0
|
|
# 如果输入为None,第二个参数为默认值，支持设置默认值。
|
|
# 如果输入为None或者空字符,第二个参数为默认值，支持设置默认值。
|
|
# 函数注释由中文/英文 双语展示
|
|
def to_bool(value: str | bool | None, default: bool = False) -> bool:
|
|
"""
|
|
此函数用于将输入值转换为布尔类型。
|
|
支持布尔值输入，也支持 'true'、'false'、'enable'、'disable'、'1'、'0' 字符串输入（忽略大小写）。
|
|
若输入为 None 或者空字符串，则返回默认值。
|
|
|
|
参数:
|
|
value (str | bool | None): 要转换为布尔值的输入，可以是字符串、布尔值或 None。
|
|
default (bool): 当输入为 None 或者空字符串时返回的默认值，默认为 False。
|
|
|
|
返回:
|
|
bool: 转换后的布尔值。
|
|
|
|
异常:
|
|
ValueError: 若输入既不是有效的布尔值，也不是支持的字符串，会抛出此异常。
|
|
|
|
This function is used to convert the input value to a boolean type.
|
|
It supports boolean input and string inputs such as 'true', 'false', 'enable', 'disable', '1', '0' (case-insensitive).
|
|
If the input is None or an empty string, the default value will be returned.
|
|
|
|
Parameters:
|
|
value (str | bool | None): The input to be converted to a boolean value, which can be a string, a boolean, or None.
|
|
default (bool): The default value to be returned when the input is None or an empty string, with a default of False.
|
|
|
|
Returns:
|
|
bool: The converted boolean value.
|
|
|
|
Exceptions:
|
|
ValueError: If the input is neither a valid boolean value nor a supported string, this exception will be thrown.
|
|
"""
|
|
if value is None or (isinstance(value, str) and value.strip() == ""):
|
|
return default
|
|
if isinstance(value, bool):
|
|
return value
|
|
if isinstance(value, str):
|
|
# 去除字符串首尾空白并转换为小写
|
|
# Remove leading and trailing whitespace from the string and convert it to lowercase
|
|
value = value.strip().lower()
|
|
if value in ['true', 'enable', '1']:
|
|
return True
|
|
elif value in ['false', 'disable', '0']:
|
|
return False
|
|
raise ValueError(f"无法将 '{value}' 转换为布尔值。 Unable to convert '{value}' to a boolean value.")
|
|
|
|
|
|
# def tensor_list_to_batch(tensor_list:list[ torch.Tensor]):
|
|
# """
|
|
# 将张量列表转换为一个批次（batch）。
|
|
|
|
# 参数:
|
|
# tensor_list (list): 张量列表。
|
|
|
|
# 返回:
|
|
# torch.Tensor: 转换后的批次张量。
|
|
# """
|
|
# return torch.stack(tensor_list, dim=0)
|
|
|
|
|
|
# def pillow_image_to_tensor(image: Image.Image):
|
|
# # 步骤 1: 将 PIL 图像对象转换为 numpy 数组
|
|
# # PIL 图像对象在内存中的存储格式不便于直接进行深度学习计算
|
|
# # 使用 np.array() 函数将其转换为 numpy 数组，数组的形状通常为 (高度, 宽度, 通道数)
|
|
# image_np = np.array(image)
|
|
|
|
# # 步骤 2: 将 numpy 数组转换为 PyTorch 张量
|
|
# # 使用 torch.tensor() 函数将 numpy 数组转换为 PyTorch 张量
|
|
# # 同时指定数据类型为 torch.float32，这是深度学习中常用的数据类型
|
|
# image_tensor = torch.tensor(image_np, dtype=torch.float32)
|
|
|
|
# # 步骤 3: 调整张量的维度
|
|
# # 在 PyTorch 中，图像张量的常见形状是 (通道数, 高度, 宽度)
|
|
# # 而 numpy 数组的形状是 (高度, 宽度, 通道数)，所以需要调整维度
|
|
# # 使用 permute() 函数对维度进行重新排列，参数 (2, 0, 1) 表示将原数组的第 2 维放到第 0 位，第 0 维放到第 1 位，第 1 维放到第 2 位
|
|
# image_tensor = image_tensor.permute(2, 0, 1)
|
|
|
|
# # 步骤 4: 归一化像素值
|
|
# # 在深度学习中，通常需要将像素值归一化到 [0, 1] 范围
|
|
# # 因为图像的像素值通常在 [0, 255] 之间，所以将每个像素值除以 255 来实现归一化
|
|
# image_tensor = image_tensor / 255
|
|
|
|
# # 步骤 5: 移除可能存在的批量维度为 1 的情况
|
|
# # 如果张量的维度是 4（通常表示 (批量大小, 通道数, 高度, 宽度)）
|
|
# # 并且批量大小为 1，那么使用 squeeze(0) 方法移除这个批量维度
|
|
# if image_tensor.ndim == 4 and image_tensor.shape[0] == 1:
|
|
# # print(f"[shatter] 当前张量维度为 {image_tensor.dim()}，自动调整为3维。")
|
|
# image_tensor = image_tensor.squeeze(0)
|
|
|
|
# # if image_tensor.dim() < 4:
|
|
# # print(f"[shatter] 当前张量维度为 {image_tensor.dim()}，自动调整为4维。")
|
|
# # # 自动调整维度
|
|
# # while image_tensor.dim() < 4:
|
|
# # image_tensor = image_tensor.unsqueeze(0)
|
|
|
|
# return image_tensor
|
|
|
|
|
|
def pillow_image_to_tensor(image: Image.Image) -> torch.Tensor:
|
|
# 步骤 1: 将 PIL 图像对象转换为 numpy 数组
|
|
# PIL 图像对象在内存中的存储格式不便于直接进行深度学习计算
|
|
# 使用 np.array() 函数将其转换为 numpy 数组，数组的形状通常为 (高度, 宽度, 通道数)
|
|
np_array = np.array(image)
|
|
# 步骤 2: 将 numpy 数组转换为 PyTorch 张量
|
|
tensor = torch.from_numpy(np_array).permute(2, 0, 1).float() / 255.0 # 转换为 (C, H, W) 形状并归一化到 [0, 1]
|
|
return tensor
|
|
|
|
def tensor_list_to_batch(tensor_list: list[torch.Tensor]) -> torch.Tensor:
|
|
"""
|
|
将张量列表转换为一个批次（batch）。
|
|
|
|
参数:
|
|
tensor_list (list): 张量列表。
|
|
|
|
返回:
|
|
torch.Tensor: 转换后的批次张量。
|
|
"""
|
|
# 检查列表是否为空
|
|
if not tensor_list:
|
|
raise ValueError("输入的张量列表为空。")
|
|
# 检查所有张量的形状是否一致
|
|
first_shape = tensor_list[0].shape
|
|
for tensor in tensor_list:
|
|
if tensor.shape != first_shape:
|
|
raise ValueError("张量列表中的所有张量形状必须一致。")
|
|
# 将张量列表堆叠成一个批次张量
|
|
batch_tensor = torch.stack(tensor_list)
|
|
return batch_tensor
|
|
|
|
|
|
# def pillow_images_to_tensor(images: list[Image.Image]) -> torch.Tensor:
|
|
# """
|
|
# 将Pillow图像列表转换为图像张量。
|
|
|
|
# 参数:
|
|
# images (list[Image.Image]): 包含Pillow图像对象的列表。
|
|
|
|
# 返回:
|
|
# torch.Tensor: 形状为 (N, C, H, W) 的图像张量，其中 N 是图像数量，C 是通道数，H 是高度，W 是宽度。
|
|
# """
|
|
# transform = transforms.Compose([
|
|
# transforms.ToTensor()
|
|
# ])
|
|
# tensors = [transform(img) for img in images]
|
|
# return torch.stack(tensors)
|
|
|
|
|
|
def tensor_common_upscale(tensor: torch.Tensor, height: int, width: int, mode: str = "bilinear", align_corners: bool = False) -> torch.Tensor:
|
|
"""
|
|
自定义的上采样函数，实现图像的上采样操作。
|
|
Custom upsampling function to perform image upsampling.
|
|
|
|
参数:
|
|
:param tensor: 输入的图像张量
|
|
The input image tensor
|
|
:param height: 目标高度
|
|
The target height
|
|
:param width: 目标宽度
|
|
The target width
|
|
:param mode: 上采样模式，默认为双线性插值
|
|
Upsampling mode, default is bilinear interpolation
|
|
:param align_corners: 是否对齐角点，默认为 False
|
|
Whether to align corners, default is False
|
|
|
|
返回:
|
|
:return: 上采样后的图像张量
|
|
The upsampled image tensor
|
|
"""
|
|
# 调整张量形状为 (B, C, H, W)
|
|
if tensor.dim() == 3:
|
|
tensor = tensor.unsqueeze(0)
|
|
# 使用 PyTorch 的 interpolate 函数进行上采样
|
|
upsampled = torch.nn.functional.interpolate(tensor, size=(height, width), mode=mode, align_corners=align_corners)
|
|
# 移除批量维度（如果只有一个样本）
|
|
if upsampled.shape[0] == 1:
|
|
upsampled = upsampled.squeeze(0)
|
|
return upsampled
|
|
|
|
|
|
def tensor_image_and_mask_from_path_or_pillow(image_input: str | Image.Image) -> tuple[torch.Tensor, torch.Tensor, bool]:
|
|
"""
|
|
处理单张图像，将其转换为 RGB 图像张量和掩码张量。
|
|
Process a single image and convert it into an RGB image tensor and a mask tensor.
|
|
|
|
参数:
|
|
:param image_input: 图像文件的路径或者 PIL.Image.Image 对象。
|
|
The path of the image file or a PIL.Image.Image object.
|
|
|
|
返回:
|
|
:return: 处理后的图像张量、掩码张量和是否有非空掩码的标志。
|
|
The processed image tensor, mask tensor, and a flag indicating whether there is a non-empty mask.
|
|
"""
|
|
if isinstance(image_input, str):
|
|
# 如果输入是字符串，当作文件路径处理
|
|
i = Image.open(image_input)
|
|
elif isinstance(image_input, Image.Image):
|
|
# 如果输入是 PIL.Image.Image 对象，直接使用
|
|
i = image_input
|
|
else:
|
|
raise ValueError("输入必须是字符串类型的文件路径或者 PIL.Image.Image 对象。")
|
|
|
|
i = ImageOps.exif_transpose(i)
|
|
# # 转换为 RGB 模式
|
|
# image = i.convert("RGB")
|
|
# 转换为 RGBA 模式
|
|
image = i.convert("RGBA")
|
|
|
|
# 将图像转换为 numpy 数组并归一化
|
|
image = np.array(image).astype(np.float32) / 255.0
|
|
# 转换为 PyTorch 张量并调整形状
|
|
image = torch.from_numpy(image)[None, ]
|
|
|
|
# 检查图像是否有透明度通道
|
|
has_non_empty_mask = False
|
|
if 'A' in i.getbands():
|
|
# 获取透明度通道并转换为 numpy 数组，然后归一化
|
|
mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
|
|
# 转换为 PyTorch 张量并取反
|
|
mask = 1. - torch.from_numpy(mask)
|
|
has_non_empty_mask = True
|
|
else:
|
|
# 如果没有透明度通道，创建全零掩码
|
|
mask = torch.zeros((64, 64), dtype=torch.float32, device="cpu")
|
|
|
|
return image, mask, has_non_empty_mask
|
|
|
|
|
|
def tensor_image_and_mask_batch_from_list(images: list[torch.Tensor], masks: list[torch.Tensor], has_non_empty_mask: bool) -> tuple[torch.Tensor, torch.Tensor, int]:
|
|
"""
|
|
根据图像和掩码列表处理并拼接图像和掩码张量。
|
|
Process and concatenate image and mask tensors based on the image and mask lists.
|
|
|
|
参数:
|
|
:param images: 图像张量列表
|
|
A list of image tensors
|
|
:param masks: 掩码张量列表
|
|
A list of mask tensors
|
|
:param has_non_empty_mask: 是否有非空掩码的标志
|
|
A flag indicating whether there is a non-empty mask
|
|
|
|
返回:
|
|
:return: 拼接后的图像张量、掩码张量和图像数量
|
|
The concatenated image tensor, mask tensor, and the number of images
|
|
"""
|
|
if len(images) == 1:
|
|
return (images[0], masks[0], 1)
|
|
elif len(images) > 1:
|
|
# 获取第一张图像的张量
|
|
image1 = images[0]
|
|
mask1 = None
|
|
|
|
# 拼接图像张量
|
|
for image2 in images[1:]:
|
|
if image1.shape[1:] != image2.shape[1:]:
|
|
# 如果图像形状不一致，进行上采样
|
|
image2 = tensor_common_upscale(image2.movedim(-1, 1), image1.shape[2], image1.shape[1], "bilinear", "center").movedim(1, -1)
|
|
image1 = torch.cat((image1, image2), dim=0)
|
|
|
|
# 拼接掩码张量
|
|
for mask2 in masks[1:]:
|
|
if has_non_empty_mask:
|
|
if image1.shape[1:3] != mask2.shape:
|
|
# 如果掩码形状与图像不一致，进行插值
|
|
mask2 = torch.nn.functional.interpolate(mask2.unsqueeze(0).unsqueeze(0), size=(image1.shape[2], image1.shape[1]), mode='bilinear', align_corners=False)
|
|
mask2 = mask2.squeeze(0)
|
|
else:
|
|
mask2 = mask2.unsqueeze(0)
|
|
else:
|
|
mask2 = mask2.unsqueeze(0)
|
|
|
|
if mask1 is None:
|
|
mask1 = mask2
|
|
else:
|
|
mask1 = torch.cat((mask1, mask2), dim=0)
|
|
|
|
return (image1, mask1, len(images))
|
|
else:
|
|
return (torch.tensor([]), torch.tensor([]), 0)
|
|
|
|
def create_shatter_effect(params: dict):
|
|
"""
|
|
创建图片破碎效果并保存序列帧
|
|
|
|
参数:
|
|
params (dict): 包含所有配置参数的字典
|
|
|
|
如果传入 shards_area_mask 碎化粒子只在该区域初始化
|
|
|
|
如果传入 save_image 则保存图片，否则不保存图片
|
|
"""
|
|
|
|
input_image_path = params.get('input_image_path')
|
|
# print(f"input_image_path type: {type(input_image_path)}, value: {input_image_path}") # 添加调试输出
|
|
# 其他代码保持不变
|
|
|
|
input_image_path = params.get('input_image_path')
|
|
output_folder = params.get('output_folder')
|
|
frame_count = params.get('frame_count', 20)
|
|
shard_size = params.get('shard_size', 30)
|
|
speed = params.get('speed', 10)
|
|
rotation = params.get('rotation', 5)
|
|
fade_start = params.get('fade_start', 0)
|
|
name_prefix = params.get('name_prefix', "frame_")
|
|
name_id_length = params.get('name_id_length', 3)
|
|
frame_start_index = params.get('frame_start_index', 1)
|
|
shard_shape = params.get('shard_shape', None)
|
|
shard_mask_image = params.get('shard_mask_image', None)
|
|
shards_area_mask = params.get('shards_area_mask', None)
|
|
save_image = to_bool(params.get('save_image', True))
|
|
show_image = to_bool(params.get('show_image', False))
|
|
|
|
frame_count,shard_size=tuple([int(item) for item in [frame_count,shard_size]])
|
|
fade_start,name_id_length,frame_start_index=tuple([int(item) for item in [fade_start,name_id_length,frame_start_index]])
|
|
speed,rotation=tuple([float(item) for item in [speed,rotation]])
|
|
# name_prefix=tuple([int(item) for item in [name_prefix]])
|
|
|
|
# 创建输出文件夹
|
|
if save_image:
|
|
os.makedirs(output_folder, exist_ok=True)
|
|
|
|
# 加载原始图片
|
|
original_img = pillow_image_from_input(input_image_path, False)
|
|
width, height = original_img.size
|
|
|
|
|
|
# 转换为RGBA模式以便处理透明度
|
|
original_img = original_img.convert("RGBA")
|
|
|
|
# 创建碎片网格
|
|
# cols = math.ceil(width / shard_size)
|
|
# rows = math.ceil(height / shard_size)
|
|
|
|
# shards_area_mask = pillow_image_from_input(shards_area_mask, is_mask=True)
|
|
# if shards_area_mask.getpixel((x, y)) != 0:
|
|
|
|
# 为每个碎片生成随机运动参数
|
|
# shards = []
|
|
# for row in range(rows):
|
|
# for col in range(cols):
|
|
# x = col * shard_size
|
|
# y = row * shard_size
|
|
# actual_width, actual_height = get_shard_dimensions(x, y, shard_size, width, height)
|
|
# if actual_width <= 0 or actual_height <= 0:
|
|
# continue
|
|
# shard_img = extract_shard_image(original_img, x, y, actual_width, actual_height)
|
|
# shard_img = apply_shard_mask(shard_img, shard_shape, shard_mask_image, actual_width, actual_height)
|
|
# shard_info = create_shard_info(shard_img, x, y, speed, rotation, frame_count, fade_start)
|
|
# shards.append(shard_info)
|
|
|
|
|
|
shards = []
|
|
|
|
if shards_area_mask is not None:
|
|
shards_area_mask = pillow_image_from_input(shards_area_mask, is_mask=True)
|
|
|
|
# 遍历原始图像的每个位置
|
|
for y in range(0, height, shard_size):
|
|
for x in range(0, width, shard_size):
|
|
|
|
# 检查遮罩图像对应位置是否为有效位置
|
|
if shards_area_mask is not None and shards_area_mask.getpixel((x, y)) == 0:
|
|
continue
|
|
|
|
# 计算碎片的实际尺寸
|
|
actual_width, actual_height = get_shard_dimensions(x, y, shard_size, width, height)
|
|
if actual_width <= 0 or actual_height <= 0:
|
|
continue
|
|
|
|
# 从原始图片中提取单个碎片图像
|
|
shard_img = extract_shard_image(original_img, x, y, actual_width, actual_height)
|
|
# 为碎片图像应用遮罩形状和遮罩图像
|
|
shard_img = apply_shard_mask(shard_img, shard_shape, shard_mask_image, actual_width, actual_height)
|
|
# 生成单个碎片的运动参数
|
|
movement_params = generate_shard_movement_params(speed, rotation, frame_count, fade_start)
|
|
# 创建单个碎片的完整信息
|
|
shard_info = create_shard_info(shard_img, x, y, speed, rotation, frame_count, fade_start)
|
|
shards.append(shard_info)
|
|
|
|
filelist:list[str]=[]
|
|
image_list = []
|
|
# 生成每一帧
|
|
for frame in range(frame_count):
|
|
# 创建透明背景
|
|
frame_img = Image.new("RGBA", (width, height), (0, 0, 0, 0))
|
|
draw = ImageDraw.Draw(frame_img)
|
|
|
|
# 绘制所有碎片
|
|
for shard in shards:
|
|
update_shard_position_and_rotation(shard, frame)
|
|
update_shard_opacity(shard, frame)
|
|
rotated_shard = rotate_and_apply_opacity(shard)
|
|
draw_x = shard['x'] - (rotated_shard.width - shard['width']) / 2
|
|
draw_y = shard['y'] - (rotated_shard.height - shard['height']) / 2
|
|
frame_img.paste(rotated_shard, (int(draw_x), int(draw_y)), rotated_shard)
|
|
|
|
|
|
# get file path
|
|
name_id = text_padd(str(frame + frame_start_index), 'padd-start', name_id_length, '0')
|
|
filepath = str(Path(f"{output_folder}/{name_prefix}{name_id}.png").as_posix())
|
|
filelist.append(filepath)
|
|
print(f"已生成帧 {frame + 1}/{frame_count}")
|
|
|
|
# 保存帧
|
|
if save_image:
|
|
frame_img.save(filepath)
|
|
|
|
if show_image:
|
|
image_list.append(frame_img)
|
|
|
|
|
|
images=[]
|
|
masks=[]
|
|
has_non_empty_mask = False
|
|
if show_image:
|
|
for img in image_list:
|
|
image,mask,has_mask = tensor_image_and_mask_from_path_or_pillow(img)
|
|
images.append(image)
|
|
masks.append(mask)
|
|
if has_mask:
|
|
has_non_empty_mask = True
|
|
images,masks,has_non_empty_mask=tensor_image_and_mask_batch_from_list(images, masks, has_non_empty_mask)
|
|
return (filelist,images,masks,has_non_empty_mask)
|
|
|