93 lines
3.8 KiB
Python
93 lines
3.8 KiB
Python
import numpy as np
|
|
from PIL import Image
|
|
import os
|
|
import glob
|
|
import torch
|
|
|
|
def images_process(images, crop_image=False, padding=0):
|
|
"""
|
|
处理图片，返回处理后的图片
|
|
:param images: 图片路径or图片张量orPIL.Image.Image
|
|
:param crop_image: 是否进行正方形裁剪，默认False
|
|
:param padding: 裁剪时额外添加的边距，默认0
|
|
:return: 处理后的图片列表
|
|
"""
|
|
images_list = []
|
|
if isinstance(images, str):
|
|
# 如果传入的是文件夹路径，则遍历文件夹中的所有图片
|
|
if os.path.isdir(images):
|
|
images_files = glob.glob(os.path.join(images, '*')) # 图片路径列表
|
|
images_list = [Image.open(image).convert("RGBA") for image in images_files] # PIL.Image图片列表
|
|
else:
|
|
images_list.append(Image.open(images).convert("RGBA"))
|
|
elif isinstance(images, torch.Tensor):
|
|
# 将传入的张量转置为(B, H, W, C)
|
|
if images.shape[1] == 3 or images.shape[1] == 4:
|
|
images = images.permute(0, 2, 3, 1) # (B, H, W, C)
|
|
|
|
print("[DEBUG]传入的是张量，张量的形状, ", images.shape)
|
|
images_tensor = [images[i:i+1] for i in range(images.shape[0])] # 将每个批次作为独立的图片
|
|
# 遍历每张贴纸
|
|
for img_tensor in images_tensor:
|
|
# 假设输入张量是 (1, H, W, C) 格式，取第一维的1（即去掉批次维度）
|
|
img_tensor = img_tensor.squeeze(0) # 去掉批次维度，得到(C, H, W)格式
|
|
# 将张量转换为PIL图像
|
|
img_np = img_tensor.cpu().numpy() # 将张量转换为numpy数组
|
|
img_np = (img_np * 255).astype(np.uint8)
|
|
# 将numpy数组转换为图像
|
|
img = Image.fromarray(img_np)
|
|
# 加入列表容器
|
|
images_list.append(img)
|
|
elif isinstance(images, list):
|
|
for image in images:
|
|
if isinstance(image, str):
|
|
image = Image.open(image).convert("RGBA")
|
|
images_list.append(image)
|
|
elif isinstance(images, Image.Image):
|
|
images_list.append(images.convert("RGBA"))
|
|
else:
|
|
print("传入的参数类型错误，", type(images))
|
|
|
|
# 在返回之前添加裁剪处理
|
|
if crop_image:
|
|
cropped_images = []
|
|
for img in images_list:
|
|
# 获取图片的alpha通道
|
|
alpha = img.split()[-1]
|
|
# 获取非透明区域的边界框
|
|
bbox = alpha.getbbox()
|
|
if bbox:
|
|
left, top, right, bottom = bbox
|
|
# 计算边界框的宽度和高度
|
|
width = right - left
|
|
height = bottom - top
|
|
# 使用最大边长构建正方形
|
|
max_side = max(width, height) + 2 * padding
|
|
|
|
# 计算原图中心点
|
|
center_x = (left + right) // 2
|
|
center_y = (top + bottom) // 2
|
|
|
|
# 计算正方形裁剪区域
|
|
new_left = center_x - max_side // 2
|
|
new_top = center_y - max_side // 2
|
|
new_right = new_left + max_side
|
|
new_bottom = new_top + max_side
|
|
|
|
# 确保裁剪区域不超出图片范围
|
|
img_width, img_height = img.size
|
|
new_left = max(0, new_left)
|
|
new_top = max(0, new_top)
|
|
new_right = min(img_width, new_right)
|
|
new_bottom = min(img_height, new_bottom)
|
|
|
|
# 裁剪图片
|
|
cropped_img = img.crop((new_left, new_top, new_right, new_bottom))
|
|
cropped_images.append(cropped_img)
|
|
else:
|
|
# 如果没有非透明区域，直接添加原图
|
|
cropped_images.append(img)
|
|
return cropped_images
|
|
|
|
return images_list
|