329 lines
8.4 KiB
Python
329 lines
8.4 KiB
Python
"""
|
||
文件处理工具模块
|
||
提供文件夹图片加载、智能命名、图片配对等功能
|
||
"""
|
||
|
||
import os
|
||
import uuid
|
||
import time
|
||
from itertools import product
|
||
from pathlib import Path
|
||
from typing import List, Tuple, Optional, NamedTuple
|
||
|
||
from PIL import Image
|
||
|
||
|
||
# 支持的图片格式
|
||
SUPPORTED_IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.gif'}
|
||
|
||
|
||
class ImageInfo(NamedTuple):
|
||
"""图片信息结构"""
|
||
image: Image.Image
|
||
filename: str # 不含扩展名的文件名
|
||
extension: str # 扩展名（如 .png）
|
||
source_path: str # 原始文件路径
|
||
|
||
|
||
def load_images_from_folder(
|
||
folder_path: str,
|
||
recursive: bool = False
|
||
) -> List[ImageInfo]:
|
||
"""
|
||
从文件夹加载所有图片
|
||
|
||
Args:
|
||
folder_path: 文件夹路径
|
||
recursive: 是否递归加载子文件夹
|
||
|
||
Returns:
|
||
ImageInfo 列表，包含图片和元数据
|
||
|
||
Raises:
|
||
ValueError: 文件夹不存在或为空
|
||
|
||
Example:
|
||
>>> images = load_images_from_folder("D:/images")
|
||
>>> for info in images:
|
||
... print(f"{info.filename}: {info.image.size}")
|
||
"""
|
||
folder_path = folder_path.strip()
|
||
|
||
if not folder_path:
|
||
return []
|
||
|
||
path = Path(folder_path)
|
||
|
||
if not path.exists():
|
||
raise ValueError(f"文件夹不存在: {folder_path}")
|
||
|
||
if not path.is_dir():
|
||
raise ValueError(f"路径不是文件夹: {folder_path}")
|
||
|
||
images = []
|
||
|
||
# 获取文件列表
|
||
if recursive:
|
||
files = list(path.rglob("*"))
|
||
else:
|
||
files = list(path.iterdir())
|
||
|
||
# 按文件名排序，确保顺序一致
|
||
files = sorted(files, key=lambda x: x.name.lower())
|
||
|
||
for file_path in files:
|
||
if not file_path.is_file():
|
||
continue
|
||
|
||
ext = file_path.suffix.lower()
|
||
if ext not in SUPPORTED_IMAGE_EXTENSIONS:
|
||
continue
|
||
|
||
try:
|
||
img = Image.open(file_path)
|
||
img.load() # 确保图片完全加载
|
||
|
||
# 转换为 RGB 模式
|
||
if img.mode != 'RGB':
|
||
img = img.convert('RGB')
|
||
|
||
images.append(ImageInfo(
|
||
image=img,
|
||
filename=file_path.stem,
|
||
extension=ext,
|
||
source_path=str(file_path)
|
||
))
|
||
except Exception as e:
|
||
print(f"警告: 无法加载图片 {file_path}: {e}")
|
||
continue
|
||
|
||
return images
|
||
|
||
|
||
def pair_images_indexed(
|
||
*image_lists: List[ImageInfo]
|
||
) -> List[Tuple[ImageInfo, ...]]:
|
||
"""
|
||
1:1 索引配对
|
||
|
||
按索引位置配对多个图片列表，以最短列表长度为准。
|
||
|
||
Args:
|
||
*image_lists: 多个 ImageInfo 列表
|
||
|
||
Returns:
|
||
配对后的元组列表
|
||
|
||
Example:
|
||
>>> list_a = [a1, a2, a3]
|
||
>>> list_b = [b1, b2, b3]
|
||
>>> pairs = pair_images_indexed(list_a, list_b)
|
||
>>> # [(a1, b1), (a2, b2), (a3, b3)]
|
||
"""
|
||
if not image_lists:
|
||
return []
|
||
|
||
# 过滤空列表
|
||
non_empty_lists = [lst for lst in image_lists if lst]
|
||
|
||
if not non_empty_lists:
|
||
return []
|
||
|
||
# 使用 zip 进行索引配对（以最短列表为准）
|
||
return list(zip(*non_empty_lists))
|
||
|
||
|
||
def pair_images_cartesian(
|
||
*image_lists: List[ImageInfo]
|
||
) -> List[Tuple[ImageInfo, ...]]:
|
||
"""
|
||
笛卡尔积配对
|
||
|
||
生成多个图片列表的所有组合。
|
||
|
||
Args:
|
||
*image_lists: 多个 ImageInfo 列表
|
||
|
||
Returns:
|
||
配对后的元组列表
|
||
|
||
Example:
|
||
>>> list_a = [a1, a2]
|
||
>>> list_b = [b1, b2]
|
||
>>> pairs = pair_images_cartesian(list_a, list_b)
|
||
>>> # [(a1, b1), (a1, b2), (a2, b1), (a2, b2)]
|
||
"""
|
||
if not image_lists:
|
||
return []
|
||
|
||
# 过滤空列表
|
||
non_empty_lists = [lst for lst in image_lists if lst]
|
||
|
||
if not non_empty_lists:
|
||
return []
|
||
|
||
# 使用 itertools.product 生成笛卡尔积
|
||
return list(product(*non_empty_lists))
|
||
|
||
|
||
def generate_output_filename(
|
||
source_images: List[ImageInfo],
|
||
batch_index: int,
|
||
output_folder: str,
|
||
extension: str = ".png",
|
||
task_id: Optional[str] = None
|
||
) -> str:
|
||
"""
|
||
生成智能输出文件名
|
||
|
||
基于源图片文件名生成输出文件名，使用任务ID和时间戳确保并发安全。
|
||
|
||
Args:
|
||
source_images: 源图片信息列表
|
||
batch_index: 批次索引（从 0 开始）
|
||
output_folder: 输出文件夹路径
|
||
extension: 输出文件扩展名
|
||
task_id: 任务唯一标识符（用于并发场景）
|
||
|
||
Returns:
|
||
完整的输出文件路径
|
||
|
||
Example:
|
||
>>> # 单图片: hello.png -> hello_task0_12345_000.png
|
||
>>> # 多图片: hello.png + ref.png -> hello_ref_task0_12345_000.png
|
||
>>> # 并发安全：每个任务有唯一的 task_id 和时间戳
|
||
"""
|
||
# 构建基础文件名
|
||
if len(source_images) == 1:
|
||
base_name = source_images[0].filename
|
||
else:
|
||
# 多个源图片，组合文件名
|
||
names = [info.filename for info in source_images]
|
||
base_name = "_".join(names)
|
||
|
||
# 确保输出文件夹存在
|
||
output_path = Path(output_folder)
|
||
output_path.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 生成唯一性标识
|
||
if task_id is None:
|
||
# 如果没有提供 task_id，使用 UUID 前8位
|
||
task_id = str(uuid.uuid4())[:8]
|
||
|
||
# 使用时间戳（毫秒级）增加唯一性
|
||
timestamp = int(time.time() * 1000) % 100000 # 精确到毫秒的后5位
|
||
|
||
# 生成文件名：基础名_任务ID_时间戳_批次索引
|
||
filename = f"{base_name}_{task_id}_{timestamp:05d}_{batch_index:03d}{extension}"
|
||
full_path = output_path / filename
|
||
|
||
# 极小概率的冲突处理
|
||
counter = 1
|
||
while full_path.exists():
|
||
filename = f"{base_name}_{task_id}_{timestamp:05d}_{batch_index:03d}_{counter}{extension}"
|
||
full_path = output_path / filename
|
||
counter += 1
|
||
|
||
return str(full_path)
|
||
|
||
|
||
def generate_batch_output_filenames(
|
||
source_images: List[ImageInfo],
|
||
count: int,
|
||
output_folder: str,
|
||
extension: str = ".png",
|
||
task_id: Optional[str] = None
|
||
) -> List[str]:
|
||
"""
|
||
批量生成输出文件名
|
||
|
||
Args:
|
||
source_images: 源图片信息列表
|
||
count: 需要生成的文件名数量
|
||
output_folder: 输出文件夹路径
|
||
extension: 输出文件扩展名
|
||
task_id: 任务唯一标识符（用于并发场景）
|
||
|
||
Returns:
|
||
输出文件路径列表
|
||
"""
|
||
filenames = []
|
||
|
||
for i in range(count):
|
||
filename = generate_output_filename(
|
||
source_images=source_images,
|
||
batch_index=i,
|
||
output_folder=output_folder,
|
||
extension=extension,
|
||
task_id=task_id
|
||
)
|
||
filenames.append(filename)
|
||
|
||
return filenames
|
||
|
||
|
||
def save_image(
|
||
image: Image.Image,
|
||
output_path: str,
|
||
quality: int = 95
|
||
) -> str:
|
||
"""
|
||
保存图片到指定路径
|
||
|
||
Args:
|
||
image: PIL Image 对象
|
||
output_path: 输出文件路径
|
||
quality: JPEG 质量（仅对 JPEG 格式有效）
|
||
|
||
Returns:
|
||
实际保存的文件路径
|
||
"""
|
||
# 确保目录存在
|
||
output_dir = Path(output_path).parent
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# 根据扩展名选择保存参数
|
||
ext = Path(output_path).suffix.lower()
|
||
|
||
if ext in {'.jpg', '.jpeg'}:
|
||
# 转换为 RGB（JPEG 不支持 alpha 通道）
|
||
if image.mode != 'RGB':
|
||
image = image.convert('RGB')
|
||
image.save(output_path, quality=quality)
|
||
elif ext == '.png':
|
||
image.save(output_path)
|
||
elif ext == '.webp':
|
||
image.save(output_path, quality=quality)
|
||
else:
|
||
image.save(output_path)
|
||
|
||
return output_path
|
||
|
||
|
||
def get_folder_image_count(folder_path: str) -> int:
|
||
"""
|
||
获取文件夹中的图片数量（不加载图片）
|
||
|
||
Args:
|
||
folder_path: 文件夹路径
|
||
|
||
Returns:
|
||
图片数量
|
||
"""
|
||
folder_path = folder_path.strip()
|
||
|
||
if not folder_path:
|
||
return 0
|
||
|
||
path = Path(folder_path)
|
||
|
||
if not path.exists() or not path.is_dir():
|
||
return 0
|
||
|
||
count = 0
|
||
for file_path in path.iterdir():
|
||
if file_path.is_file() and file_path.suffix.lower() in SUPPORTED_IMAGE_EXTENSIONS:
|
||
count += 1
|
||
|
||
return count
|