""" 文件处理工具模块 提供文件夹图片加载、智能命名、图片配对等功能 """ 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