import argparse import logging from pathlib import Path import random import sys import json # Added for reading metadata try: from PIL import Image except ImportError: print("错误: 未找到所需的库 PIL (Pillow) 和 imagehash。") print("请先安装它们: pip install Pillow imagehash") sys.exit(1) # --- 配置 --- # 日志设置 logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s', datefmt='%Y-%m-%d %H:%M:%S' ) logger = logging.getLogger(__name__) # 支持的图片文件扩展名 (小写) IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.webp'} # 每组选取的图片数量 IMAGES_PER_GROUP = 3 # --- 核心功能 --- def find_image_files(directory: Path) -> list[Path]: # Changed to non-recursive find to align with poster_notes_creator logic """Finds image files directly within the specified directory (non-recursive).""" if not directory.is_dir(): logger.error(f"错误：目录不存在或不是一个有效的目录 -> {directory}") return [] logger.info(f"开始在目录中扫描图片 (非递归): {directory}") image_files = [] try: # Use iterdir for non-recursive listing for item in directory.iterdir(): if item.is_file() and item.suffix.lower() in IMAGE_EXTENSIONS: image_files.append(item) except Exception as e: logger.error(f"扫描目录时出错: {e}") return [] logger.info(f"在 {directory} 中找到 {len(image_files)} 个图片文件。") return image_files def calculate_image_hash(image_path: Path) -> str | None: """计算图片的感知哈希 (phash) 作为数字指纹。""" try: # 使用 Pillow 打开图片 with Image.open(image_path) as img: # phash 通常在灰度图上效果更好，且能处理不同模式（如 RGBA, P） # 如果图片无法转换或打开，会抛出异常 hash_val = imagehash.phash(img.convert('L')) # 返回哈希值的字符串形式 return str(hash_val) except Exception as e: logger.warning(f"无法识别的图片格式，跳过哈希计算: {image_path}") return None except Exception as e: # 捕获其他可能的错误，如文件损坏、权限问题等 logger.warning(f"无法为图片计算哈希值 {image_path}: {e}") return None def select_image_groups( image_paths: list[Path], num_groups: int, metadata_path: Path # Added metadata path argument ) -> list[list[dict]]: """从图片列表中随机选择指定数量的图片组，排除元数据中指定的图片。""" # --- Read Metadata and Filter Used Images --- used_image_filenames = set() if metadata_path and metadata_path.is_file(): try: with open(metadata_path, 'r', encoding='utf-8') as f: metadata = json.load(f) # Assuming the key for used images is 'collage_images' # Adjust the key if your metadata structure is different used_list = metadata.get('collage_images', []) if isinstance(used_list, list): used_image_filenames = set(used_list) logger.info(f"从元数据 {metadata_path.name} 中加载了 {len(used_image_filenames)} 个已使用的图片文件名进行排除。") else: logger.warning(f"元数据文件 {metadata_path.name} 中的 'collage_images' 键不是一个列表，无法排除图片。") except json.JSONDecodeError: logger.error(f"无法解析元数据文件 (JSON格式错误): {metadata_path}. 将不排除任何图片。") except Exception as e: logger.error(f"读取元数据文件时出错 {metadata_path}: {e}. 将不排除任何图片。") elif metadata_path: logger.warning(f"指定的元数据文件不存在或不是文件: {metadata_path}. 将不排除任何图片。") else: logger.info("未提供元数据文件，将不排除任何图片。") # Filter the initial list based on used filenames available_image_paths = [p for p in image_paths if p.name not in used_image_filenames] logger.info(f"总共找到 {len(image_paths)} 张图片, 排除 {len(used_image_filenames)} 张已用图片后, 剩余可用图片: {len(available_image_paths)} 张。") # --- End Metadata Reading and Filtering --- total_available_images = len(available_image_paths) if total_available_images < IMAGES_PER_GROUP: logger.error(f"可用图片总数 ({total_available_images}) 少于每组所需的 {IMAGES_PER_GROUP} 张。无法进行选择。") return [] selected_groups_data = [] logger.info(f"开始从 {total_available_images} 张可用图片中选择 {num_groups} 组，每组 {IMAGES_PER_GROUP} 张...") # Use a copy of the available paths list for sampling if groups might overlap significantly # For now, sample directly from the available list current_available_paths = list(available_image_paths) # Make a copy if we need to ensure groups don't deplete the pool for later groups in one run for i in range(num_groups): logger.info(f"--- 正在选择第 {i + 1} / {num_groups} 组 ---") # Check if enough images remain for the current group if len(current_available_paths) < IMAGES_PER_GROUP: logger.warning(f"剩余可用图片 ({len(current_available_paths)}) 不足以组成第 {i+1} 组 (需要 {IMAGES_PER_GROUP} 张)。停止选择。") break # Stop if not enough images left for a full group try: # Sample from the *currently available* paths selected_paths_for_group = random.sample(current_available_paths, IMAGES_PER_GROUP) # Optional: Remove selected images from the pool if groups must be distinct # (Commented out for now, allowing images to appear in multiple groups) # for path in selected_paths_for_group: # current_available_paths.remove(path) group_data = [] for img_path in selected_paths_for_group: img_hash = calculate_image_hash(img_path) group_data.append({ "path": str(img_path.resolve()), "hash": img_hash }) logger.debug(f" 已选定: {img_path} (哈希: {img_hash if img_hash else '无法计算'})") selected_groups_data.append(group_data) except ValueError: # This might happen if the pool size becomes less than IMAGES_PER_GROUP between check and sample (unlikely here) logger.error(f"选择第 {i+1} 组时发生 ValueError (剩余图片不足？)。已选 {len(current_available_paths)} 张可用。") break except Exception as e: logger.error(f"选择第 {i+1} 组时发生意外错误: {e}") continue # Skip this group logger.info(f"图片组选择完成。共选出 {len(selected_groups_data)} 组。") return selected_groups_data def print_results(selected_groups: list[list[dict]]): """格式化并打印选择结果。""" if not selected_groups: print("\n未能选出任何图片组。") return print("\n========== 图片选择结果 ==========") for i, group in enumerate(selected_groups): print(f"\n--- 组 {i + 1} ---") if not group: print(" (此组选择失败)") continue for item in group: print(f" 图片路径: {item['path']}") print(f" 数字指纹 (phash): {item['hash'] if item['hash'] else 'N/A'}") print("\n===================================") # --- 主程序入口 --- if __name__ == "__main__": # 设置命令行参数解析 parser = argparse.ArgumentParser( description=f"从指定目录随机选择数组图片（每组 {IMAGES_PER_GROUP} 张），并计算数字指纹。" ) parser.add_argument( "--image_dir", type=str, required=True, help="包含图片的源目录路径。" ) parser.add_argument( "--num_groups", type=int, required=True, help="要选择的图片组数量。" ) parser.add_argument( "--metadata_path", type=str, required=True, # Make it required as per user request help="包含已使用图片列表的 JSON 元数据文件路径。假设键为 'collage_images'。" ) # 解析参数 args = parser.parse_args() # 输入验证 source_directory = Path(args.image_dir) number_of_groups = args.num_groups metadata_file_path = Path(args.metadata_path) if number_of_groups <= 0: logger.error("选择的组数必须是一个正整数。") if not metadata_file_path.is_file(): logger.error(f"指定的元数据文件不存在或不是一个文件: {metadata_file_path}") sys.exit(1) # 1. 查找图片 (Non-recursive) all_images = find_image_files(source_directory) if not all_images: logger.info("在指定目录中未找到任何符合条件的图片文件。") # Changed from error to info sys.exit(0) # Exit gracefully if no images found # 2. 选择图片组 (Now includes filtering) final_selected_groups = select_image_groups(all_images, number_of_groups, metadata_file_path) # 3. 打印结果 print_results(final_selected_groups) logger.info("脚本执行完毕。")