699 lines
29 KiB
Python
699 lines
29 KiB
Python
import os
|
||
import random
|
||
import logging
|
||
import json
|
||
import argparse
|
||
from typing import List, Dict, Any, Tuple, Optional
|
||
from PIL import Image, ImageEnhance, ImageFilter, ImageOps # 添加 PIL 导入
|
||
import colorsys # 添加 colorsys 导入
|
||
|
||
# Setup logging
|
||
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
|
||
logger = logging.getLogger(__name__)
|
||
|
||
def select_image_batch(
|
||
available_images: List[str],
|
||
num_images_to_select: int,
|
||
batch_seed: int
|
||
) -> List[str]:
|
||
"""
|
||
使用给定种子从可用图像列表中随机选择指定数量的图像。
|
||
|
||
Args:
|
||
available_images: 可用图像文件名的列表。
|
||
num_images_to_select: 此批次要选择的图像数量。
|
||
batch_seed: 此批次选择使用的随机种子。
|
||
|
||
Returns:
|
||
List[str]: 该批次选定的图像文件名列表。
|
||
"""
|
||
if not available_images:
|
||
logger.warning("没有可供选择的图像。")
|
||
return []
|
||
|
||
# 如果可用图像少于请求数量，发出警告但选择所有可用图像
|
||
if len(available_images) < num_images_to_select:
|
||
logger.warning(
|
||
f"可用图像数量 ({len(available_images)}) 少于请求的数量 ({num_images_to_select})，"
|
||
f"将选择所有可用图像。"
|
||
)
|
||
# 即使选择所有图像，也使用 sample 进行随机排序
|
||
random.seed(batch_seed)
|
||
selected = random.sample(available_images, len(available_images))
|
||
random.seed() # 重置种子以避免影响后续操作
|
||
return selected
|
||
else:
|
||
logger.info(f"使用种子 {batch_seed} 从 {len(available_images)} 张可用图像中选择 {num_images_to_select} 张。")
|
||
random.seed(batch_seed)
|
||
selected = random.sample(available_images, num_images_to_select)
|
||
random.seed() # 重置种子
|
||
return selected
|
||
|
||
# --- 从 poster_notes_creator.py 复制的图像处理函数 --- START ---
|
||
|
||
def _adjust_brightness(image: Image.Image, factor: float) -> Image.Image:
|
||
"""调整图像亮度"""
|
||
if factor == 1.0:
|
||
return image
|
||
|
||
try:
|
||
enhancer = ImageEnhance.Brightness(image)
|
||
return enhancer.enhance(factor)
|
||
except ImportError:
|
||
logger.warning("ImageEnhance.Brightness not available, using manual adjustment.")
|
||
data = list(image.getdata())
|
||
new_data = []
|
||
|
||
for pixel in data:
|
||
r, g, b = pixel[:3]
|
||
r = min(255, max(0, int(r * factor)))
|
||
g = min(255, max(0, int(g * factor)))
|
||
b = min(255, max(0, int(b * factor)))
|
||
|
||
if len(pixel) > 3: # 如果有alpha通道
|
||
new_data.append((r, g, b, pixel[3]))
|
||
else:
|
||
new_data.append((r, g, b))
|
||
|
||
result = Image.new(image.mode, image.size)
|
||
result.putdata(new_data)
|
||
return result
|
||
|
||
def _adjust_contrast(image: Image.Image, factor: float) -> Image.Image:
|
||
"""调整图像对比度"""
|
||
if factor == 1.0:
|
||
return image
|
||
|
||
try:
|
||
enhancer = ImageEnhance.Contrast(image)
|
||
return enhancer.enhance(factor)
|
||
except ImportError:
|
||
logger.warning("ImageEnhance.Contrast not available, using manual adjustment.")
|
||
data = list(image.getdata())
|
||
new_data = []
|
||
|
||
# 计算平均亮度
|
||
avg_r, avg_g, avg_b = 0, 0, 0
|
||
count = 0
|
||
|
||
for pixel in data:
|
||
r, g, b = pixel[:3]
|
||
avg_r += r
|
||
avg_g += g
|
||
avg_b += b
|
||
count += 1
|
||
|
||
if count > 0:
|
||
avg_r //= count
|
||
avg_g //= count
|
||
avg_b //= count
|
||
else:
|
||
return image # 无法计算平均值
|
||
|
||
# 调整对比度
|
||
for pixel in data:
|
||
r, g, b = pixel[:3]
|
||
r = min(255, max(0, int(avg_r + (r - avg_r) * factor)))
|
||
g = min(255, max(0, int(avg_g + (g - avg_g) * factor)))
|
||
b = min(255, max(0, int(avg_b + (b - avg_b) * factor)))
|
||
|
||
if len(pixel) > 3: # 如果有alpha通道
|
||
new_data.append((r, g, b, pixel[3]))
|
||
else:
|
||
new_data.append((r, g, b))
|
||
|
||
result = Image.new(image.mode, image.size)
|
||
result.putdata(new_data)
|
||
return result
|
||
|
||
def _adjust_saturation(image: Image.Image, factor: float) -> Image.Image:
|
||
"""调整图像饱和度"""
|
||
if factor == 1.0:
|
||
return image
|
||
|
||
try:
|
||
enhancer = ImageEnhance.Color(image)
|
||
return enhancer.enhance(factor)
|
||
except ImportError:
|
||
logger.warning("ImageEnhance.Color not available, using manual adjustment.")
|
||
data = list(image.getdata())
|
||
new_data = []
|
||
|
||
for pixel in data:
|
||
r, g, b = pixel[:3]
|
||
# 计算灰度值
|
||
gray = int(0.299 * r + 0.587 * g + 0.114 * b) # 更精确的灰度计算
|
||
# 调整饱和度
|
||
r = min(255, max(0, int(gray + (r - gray) * factor)))
|
||
g = min(255, max(0, int(gray + (g - gray) * factor)))
|
||
b = min(255, max(0, int(gray + (b - gray) * factor)))
|
||
|
||
if len(pixel) > 3: # 如果有alpha通道
|
||
new_data.append((r, g, b, pixel[3]))
|
||
else:
|
||
new_data.append((r, g, b))
|
||
|
||
result = Image.new(image.mode, image.size)
|
||
result.putdata(new_data)
|
||
return result
|
||
|
||
def _adjust_hue(image: Image.Image, shift: float) -> Image.Image:
|
||
"""调整图像色相"""
|
||
if shift == 0.0:
|
||
return image
|
||
|
||
try:
|
||
# 获取像素数据
|
||
image = image.convert('RGB') # 确保是RGB模式
|
||
data = list(image.getdata())
|
||
new_data = []
|
||
|
||
for pixel in data:
|
||
r, g, b = pixel
|
||
# 转换为HSV
|
||
h, s, v = colorsys.rgb_to_hsv(r/255.0, g/255.0, b/255.0)
|
||
# 调整色相 (H 是 0-1 的值)
|
||
h = (h + shift) % 1.0
|
||
# 转回RGB
|
||
r_new, g_new, b_new = colorsys.hsv_to_rgb(h, s, v)
|
||
r_new = int(r_new * 255)
|
||
g_new = int(g_new * 255)
|
||
b_new = int(b_new * 255)
|
||
|
||
new_data.append((r_new, g_new, b_new))
|
||
|
||
result = Image.new(image.mode, image.size)
|
||
result.putdata(new_data)
|
||
return result
|
||
except Exception as e:
|
||
logger.warning(f"Error adjusting hue: {e}. Returning original image.")
|
||
return image
|
||
|
||
def _add_noise(image: Image.Image, intensity: float = 0.02) -> Image.Image:
|
||
"""添加微弱噪点，intensity控制噪点强度(0-1)"""
|
||
if intensity <= 0:
|
||
return image
|
||
|
||
try:
|
||
image = image.convert('RGB') # 确保是RGB模式
|
||
data = list(image.getdata())
|
||
new_data = []
|
||
noise_limit = int(intensity * 255)
|
||
|
||
for pixel in data:
|
||
r, g, b = pixel
|
||
# 添加随机噪点
|
||
noise_r = random.randint(-noise_limit, noise_limit)
|
||
noise_g = random.randint(-noise_limit, noise_limit)
|
||
noise_b = random.randint(-noise_limit, noise_limit)
|
||
|
||
r = max(0, min(255, r + noise_r))
|
||
g = max(0, min(255, g + noise_g))
|
||
b = max(0, min(255, b + noise_b))
|
||
|
||
new_data.append((r, g, b))
|
||
|
||
result = Image.new(image.mode, image.size)
|
||
result.putdata(new_data)
|
||
return result
|
||
except Exception as e:
|
||
logger.warning(f"Error adding noise: {e}. Returning original image.")
|
||
return image
|
||
|
||
def _add_border_and_crop(image: Image.Image, border_size: int) -> Image.Image:
|
||
"""添加边框然后裁剪回原尺寸，用于改变边缘像素"""
|
||
if border_size <= 0:
|
||
return image
|
||
|
||
try:
|
||
width, height = image.size
|
||
image = image.convert('RGB') # 确保是RGB模式
|
||
|
||
# 创建略大的画布
|
||
border_color = (
|
||
random.randint(0, 10),
|
||
random.randint(0, 10),
|
||
random.randint(0, 10)
|
||
)
|
||
bordered = Image.new(image.mode, (width + border_size*2, height + border_size*2), border_color)
|
||
bordered.paste(image, (border_size, border_size))
|
||
|
||
# 随机裁剪回原尺寸
|
||
offset_x = random.randint(0, border_size*2)
|
||
offset_y = random.randint(0, border_size*2)
|
||
result = bordered.crop((offset_x, offset_y, offset_x + width, offset_y + height))
|
||
return result
|
||
except Exception as e:
|
||
logger.warning(f"Error adding border and cropping: {e}. Returning original image.")
|
||
return image
|
||
|
||
def _slight_sharpen(image: Image.Image) -> Image.Image:
|
||
"""轻微锐化图像"""
|
||
try:
|
||
enhancer = ImageEnhance.Sharpness(image)
|
||
return enhancer.enhance(1.2) # 轻微锐化，1.0是原始锐度
|
||
except ImportError:
|
||
logger.warning("ImageEnhance.Sharpness not available. Skipping sharpen.")
|
||
return image
|
||
except Exception as e:
|
||
logger.warning(f"Error applying sharpen: {e}. Returning original image.")
|
||
return image
|
||
|
||
def _slight_blur(image: Image.Image) -> Image.Image:
|
||
"""轻微模糊图像"""
|
||
try:
|
||
return image.filter(ImageFilter.GaussianBlur(radius=0.5))
|
||
except ImportError:
|
||
logger.warning("ImageFilter.GaussianBlur not available. Skipping blur.")
|
||
return image
|
||
except Exception as e:
|
||
logger.warning(f"Error applying blur: {e}. Returning original image.")
|
||
return image
|
||
|
||
def process_image_to_aspect_ratio(
|
||
image: Image.Image,
|
||
target_ratio: Tuple[int, int],
|
||
add_variation: bool = True,
|
||
seed: int = None,
|
||
variation_strength: str = "medium",
|
||
extra_effects: bool = True
|
||
) -> Optional[Image.Image]:
|
||
"""
|
||
处理图像到指定的宽高比，并添加微小变化
|
||
Args:
|
||
image: 原始图像
|
||
target_ratio: 目标宽高比，如(3, 4)
|
||
add_variation: 是否添加微小变化
|
||
seed: 随机种子
|
||
variation_strength: 微调强度 ("low", "medium", "high")
|
||
extra_effects: 是否添加额外效果
|
||
Returns:
|
||
Image.Image: 处理后的图像, 或在失败时返回 None
|
||
"""
|
||
try:
|
||
original_mode = image.mode
|
||
# 如果指定了种子，设置随机种子
|
||
if seed is not None:
|
||
random.seed(seed)
|
||
|
||
# 根据微调强度设置参数范围
|
||
if variation_strength == "low":
|
||
brightness_range = (-0.03, 0.03)
|
||
contrast_range = (-0.03, 0.03)
|
||
saturation_range = (-0.03, 0.03)
|
||
hue_range = (-0.01, 0.01)
|
||
max_crop_px = 3
|
||
max_rotation = 0.5
|
||
noise_intensity = 0.01
|
||
border_size_range = (0, 2)
|
||
elif variation_strength == "high":
|
||
brightness_range = (-0.08, 0.08)
|
||
contrast_range = (-0.08, 0.08)
|
||
saturation_range = (-0.08, 0.08)
|
||
hue_range = (-0.02, 0.02)
|
||
max_crop_px = 8
|
||
max_rotation = 2.0
|
||
noise_intensity = 0.03
|
||
border_size_range = (0, 4)
|
||
else: # medium (默认)
|
||
brightness_range = (-0.05, 0.05)
|
||
contrast_range = (-0.05, 0.05)
|
||
saturation_range = (-0.05, 0.05)
|
||
hue_range = (-0.015, 0.015)
|
||
max_crop_px = 5
|
||
max_rotation = 1.0
|
||
noise_intensity = 0.02
|
||
border_size_range = (0, 3)
|
||
|
||
width, height = image.size
|
||
current_ratio = width / height
|
||
target_ratio_value = target_ratio[0] / target_ratio[1]
|
||
|
||
# 调整尺寸策略：使目标尺寸大约为 900x1200 或 1200x900 (3:4)
|
||
# 您可以根据需要修改这里的基准尺寸
|
||
base_width = 900
|
||
base_height = 1200
|
||
if target_ratio_value > 1: # 宽幅图，交换基准尺寸
|
||
base_width, base_height = base_height, base_width
|
||
|
||
if current_ratio > target_ratio_value: # 图片较宽
|
||
new_height = base_height
|
||
new_width = int(new_height * current_ratio)
|
||
else: # 图片较高
|
||
new_width = base_width
|
||
new_height = int(new_width / current_ratio)
|
||
|
||
# 避免尺寸为0
|
||
new_width = max(1, new_width)
|
||
new_height = max(1, new_height)
|
||
|
||
resized_image = image.resize((new_width, new_height), Image.LANCZOS)
|
||
|
||
# 计算裁剪区域
|
||
resized_width, resized_height = resized_image.size
|
||
if resized_width / resized_height > target_ratio_value:
|
||
crop_width = int(resized_height * target_ratio_value)
|
||
crop_height = resized_height
|
||
# 微小偏移
|
||
max_offset_x = max(1, min(20, (resized_width - crop_width) // 5)) if add_variation else 0
|
||
offset_x = random.randint(-max_offset_x, max_offset_x) if add_variation else 0
|
||
crop_x1 = max(0, min((resized_width - crop_width) // 2 + offset_x, resized_width - crop_width))
|
||
crop_y1 = 0
|
||
else:
|
||
crop_height = int(resized_width / target_ratio_value)
|
||
crop_width = resized_width
|
||
# 微小偏移
|
||
max_offset_y = max(1, min(20, (resized_height - crop_height) // 5)) if add_variation else 0
|
||
offset_y = random.randint(-max_offset_y, max_offset_y) if add_variation else 0
|
||
crop_y1 = max(0, min((resized_height - crop_height) // 2 + offset_y, resized_height - crop_height))
|
||
crop_x1 = 0
|
||
|
||
crop_x2 = crop_x1 + crop_width
|
||
crop_y2 = crop_y1 + crop_height
|
||
|
||
result = resized_image.crop((crop_x1, crop_y1, crop_x2, crop_y2))
|
||
|
||
if add_variation:
|
||
# 确保处理时为 RGB 或 RGBA
|
||
if result.mode not in ('RGB', 'RGBA'):
|
||
processed_image = result.convert('RGB')
|
||
else:
|
||
processed_image = result.copy()
|
||
|
||
# 1. 亮度
|
||
brightness_factor = 1.0 + random.uniform(*brightness_range)
|
||
processed_image = _adjust_brightness(processed_image, brightness_factor)
|
||
|
||
# 2. 对比度
|
||
contrast_factor = 1.0 + random.uniform(*contrast_range)
|
||
processed_image = _adjust_contrast(processed_image, contrast_factor)
|
||
|
||
# 3. 饱和度
|
||
saturation_factor = 1.0 + random.uniform(*saturation_range)
|
||
processed_image = _adjust_saturation(processed_image, saturation_factor)
|
||
|
||
# 4. 微小裁剪
|
||
crop_px = random.randint(0, max_crop_px)
|
||
if crop_px > 0:
|
||
w, h = processed_image.size
|
||
if w > 2*crop_px and h > 2*crop_px:
|
||
processed_image = processed_image.crop((crop_px, crop_px, w-crop_px, h-crop_px))
|
||
processed_image = processed_image.resize((w, h), Image.LANCZOS)
|
||
|
||
# 5. 微小旋转
|
||
rotation_angle = random.uniform(-max_rotation, max_rotation)
|
||
try:
|
||
# 使用 expand=False 保持尺寸，填充背景可能为黑色，需要处理
|
||
processed_image = processed_image.rotate(rotation_angle, resample=Image.BICUBIC, expand=False)
|
||
except ValueError:
|
||
logger.warning(f"Skipping rotation due to potential transparency issues.")
|
||
|
||
# 6. 额外效果
|
||
if extra_effects:
|
||
# 6.1 噪点
|
||
processed_image = _add_noise(processed_image, intensity=noise_intensity)
|
||
# 6.2 色相
|
||
processed_image = _adjust_hue(processed_image, shift=random.uniform(*hue_range))
|
||
# 6.3 边缘微调
|
||
border_size = random.randint(*border_size_range)
|
||
if border_size > 0:
|
||
processed_image = _add_border_and_crop(processed_image, border_size)
|
||
# 6.4 锐化/模糊
|
||
if random.random() > 0.5:
|
||
processed_image = _slight_sharpen(processed_image)
|
||
else:
|
||
processed_image = _slight_blur(processed_image)
|
||
|
||
# 如果原始是 RGBA，尝试恢复 alpha 通道
|
||
# 注意：旋转等操作可能使alpha处理复杂，这里简单处理
|
||
if original_mode == 'RGBA' and processed_image.mode == 'RGB':
|
||
try:
|
||
alpha = result.getchannel('A')
|
||
processed_image.putalpha(alpha)
|
||
except (ValueError, IndexError):
|
||
logger.warning("Could not restore original alpha channel after processing.")
|
||
elif original_mode != processed_image.mode and original_mode != 'P': # 避免转换回索引调色板
|
||
try:
|
||
processed_image = processed_image.convert(original_mode)
|
||
except ValueError:
|
||
logger.warning(f"Could not convert processed image back to original mode {original_mode}.")
|
||
|
||
final_image = processed_image
|
||
|
||
else:
|
||
final_image = result
|
||
|
||
# 重置随机种子
|
||
if seed is not None:
|
||
random.seed()
|
||
|
||
return final_image
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error processing image: {e}")
|
||
# 重置随机种子以防万一
|
||
if seed is not None:
|
||
random.seed()
|
||
return None
|
||
|
||
# --- 从 poster_notes_creator.py 复制的图像处理函数 --- END ---
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="根据指定的目录和批次数，生成指定数量的选图批次，并可选择处理图像。")
|
||
parser.add_argument(
|
||
"--source-dir",
|
||
required=True,
|
||
help="包含源图像的目录路径。"
|
||
)
|
||
parser.add_argument(
|
||
"--num-batches",
|
||
type=int,
|
||
required=True,
|
||
help="要生成的图像批次数。"
|
||
)
|
||
parser.add_argument(
|
||
"--images-per-batch",
|
||
type=int,
|
||
required=True,
|
||
help="每批要选择的图像数量。"
|
||
)
|
||
parser.add_argument(
|
||
"--base-seed",
|
||
type=int,
|
||
default=42,
|
||
help="用于生成批次选择的随机种子基数 (默认为 42)。"
|
||
)
|
||
parser.add_argument(
|
||
"--output-file",
|
||
default="selected_batches.json",
|
||
help="用于保存选定批次信息的JSON文件路径 (默认为 selected_batches.json)。"
|
||
)
|
||
|
||
# 添加图像处理相关参数
|
||
parser.add_argument(
|
||
"--process-images",
|
||
action='store_true', # 设为开关参数
|
||
help="如果指定，则对选定的图像进行处理。"
|
||
)
|
||
parser.add_argument(
|
||
"--target-ratio",
|
||
default="3:4",
|
||
help="处理图像的目标宽高比 (格式 W:H, 例如 3:4, 16:9)。默认为 3:4。"
|
||
)
|
||
parser.add_argument(
|
||
"--variation-strength",
|
||
choices=["low", "medium", "high"],
|
||
default="medium",
|
||
help="图像处理的微调强度。默认为 medium。"
|
||
)
|
||
parser.add_argument(
|
||
"--extra-effects",
|
||
action=argparse.BooleanOptionalAction, # 提供 --extra-effects / --no-extra-effects
|
||
default=True,
|
||
help="是否在图像处理中添加额外效果 (噪点、色相等)。默认启用。"
|
||
)
|
||
parser.add_argument(
|
||
"--output-processed-dir",
|
||
default="processed_batches",
|
||
help="用于保存处理后图像的目录 (默认为 processed_batches)。仅当 --process-images 被指定时使用。"
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 解析 target_ratio
|
||
try:
|
||
ratio_parts = args.target_ratio.split(':')
|
||
if len(ratio_parts) != 2:
|
||
raise ValueError("Ratio must be in W:H format")
|
||
target_ratio_tuple = (int(ratio_parts[0]), int(ratio_parts[1]))
|
||
if target_ratio_tuple[0] <= 0 or target_ratio_tuple[1] <= 0:
|
||
raise ValueError("Ratio dimensions must be positive")
|
||
except ValueError as e:
|
||
logger.error(f"无效的目标宽高比 '{args.target_ratio}': {e}")
|
||
return
|
||
|
||
# 检查源目录是否存在
|
||
if not os.path.isdir(args.source_dir):
|
||
logger.error(f"源目录不存在: {args.source_dir}")
|
||
return
|
||
|
||
# 列出源目录中的所有图像文件
|
||
image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff', '.webp') # 支持更多格式
|
||
try:
|
||
all_files = os.listdir(args.source_dir)
|
||
# 使用排序确保每次运行时文件列表顺序一致（在随机选择前）
|
||
available_images = sorted([
|
||
f for f in all_files
|
||
if os.path.isfile(os.path.join(args.source_dir, f)) and
|
||
f.lower().endswith(image_extensions)
|
||
])
|
||
except OSError as e:
|
||
logger.error(f"无法读取目录 {args.source_dir}: {e}")
|
||
return
|
||
|
||
if not available_images:
|
||
logger.error(f"在源目录 {args.source_dir} 中没有找到支持的图像文件。")
|
||
return
|
||
|
||
logger.info(f"在 {args.source_dir} 中找到 {len(available_images)} 张可用图像。")
|
||
|
||
# 检查是否有足够的图像（如果需要无重复选择 across batches）
|
||
total_images_needed_if_no_repeat = args.num_batches * args.images_per_batch
|
||
if total_images_needed_if_no_repeat > len(available_images):
|
||
logger.warning(
|
||
f"请求的总图像数 ({total_images_needed_if_no_repeat}) 大于可用图像数 ({len(available_images)})。"
|
||
f"不同批次之间可能会选择到相同的图像。"
|
||
)
|
||
# 注意: 当前实现是每个批次独立地从完整的可用图像列表中抽样。
|
||
|
||
all_batches: Dict[str, List[str]] = {}
|
||
|
||
# 生成并选择每个批次
|
||
for i in range(args.num_batches):
|
||
batch_index = i + 1
|
||
# 为每个批次生成确定性种子
|
||
batch_seed = args.base_seed + i
|
||
logger.info(f"--- 开始处理批次 {batch_index}/{args.num_batches} (种子: {batch_seed}) ---")
|
||
|
||
selected_for_batch = select_image_batch(
|
||
available_images=available_images,
|
||
num_images_to_select=args.images_per_batch,
|
||
batch_seed=batch_seed
|
||
)
|
||
|
||
if selected_for_batch:
|
||
batch_name = f"batch_{batch_index}"
|
||
all_batches[batch_name] = selected_for_batch
|
||
# 打印部分选择结果以供预览
|
||
preview = ', '.join(selected_for_batch[:5])
|
||
if len(selected_for_batch) > 5:
|
||
preview += '...'
|
||
logger.info(f"批次 {batch_index} 选择完成，选择了 {len(selected_for_batch)} 张图像: [{preview}]")
|
||
|
||
# --- 新增：处理选定的图像 --- START ---
|
||
if args.process_images and selected_for_batch:
|
||
logger.info(f"开始处理批次 {batch_index} 的 {len(selected_for_batch)} 张图像...")
|
||
processed_output_dir = args.output_processed_dir
|
||
|
||
# 确保处理输出目录存在
|
||
try:
|
||
if not os.path.exists(processed_output_dir):
|
||
os.makedirs(processed_output_dir)
|
||
logger.info(f"创建处理后图像的输出目录: {processed_output_dir}")
|
||
except OSError as e:
|
||
logger.error(f"无法创建处理输出目录 {processed_output_dir}: {e}. 跳过处理。")
|
||
continue # 跳过这个批次的处理
|
||
|
||
processed_count = 0
|
||
for img_idx, original_filename in enumerate(selected_for_batch):
|
||
input_image_path = os.path.join(args.source_dir, original_filename)
|
||
# 构建输出文件名
|
||
base, ext = os.path.splitext(original_filename)
|
||
# 保持原始扩展名，或者统一为 jpg？这里先保持
|
||
output_filename = f"batch_{batch_index}_processed_{base}{ext}"
|
||
output_image_path = os.path.join(processed_output_dir, output_filename)
|
||
|
||
# 为每个图像使用不同的微调种子
|
||
variation_seed = batch_seed + img_idx
|
||
|
||
try:
|
||
with Image.open(input_image_path) as img:
|
||
logger.debug(f"处理图像: {original_filename} -> {output_filename} (种子: {variation_seed})")
|
||
processed_img = process_image_to_aspect_ratio(
|
||
image=img,
|
||
target_ratio=target_ratio_tuple,
|
||
add_variation=True, # 始终添加变化，由强度控制
|
||
seed=variation_seed,
|
||
variation_strength=args.variation_strength,
|
||
extra_effects=args.extra_effects
|
||
)
|
||
|
||
if processed_img:
|
||
# 保存处理后的图像
|
||
# 尝试保存为原始格式，如果失败则回退到 JPEG
|
||
try:
|
||
save_format = Image.registered_extensions().get(ext.lower())
|
||
if not save_format or save_format == 'JPEG': # 对 JPEG 设置质量
|
||
processed_img.save(output_image_path, quality=95, format='JPEG')
|
||
elif save_format == 'PNG':
|
||
processed_img.save(output_image_path, optimize=True, format='PNG')
|
||
else:
|
||
processed_img.save(output_image_path, format=save_format)
|
||
except (KeyError, ValueError, IOError, OSError) as save_err:
|
||
logger.warning(f"无法按原始格式 {ext} 保存 {output_filename} ({save_err}). 尝试保存为 JPEG.")
|
||
output_filename = f"batch_{batch_index}_processed_{base}.jpg"
|
||
output_image_path = os.path.join(processed_output_dir, output_filename)
|
||
try:
|
||
# 确保是RGB模式才能保存为JPEG
|
||
if processed_img.mode == 'RGBA':
|
||
rgb_img = Image.new("RGB", processed_img.size, (255, 255, 255))
|
||
rgb_img.paste(processed_img, mask=processed_img.split()[3]) # 应用alpha蒙版
|
||
rgb_img.save(output_image_path, quality=95, format='JPEG')
|
||
elif processed_img.mode != 'RGB':
|
||
processed_img.convert('RGB').save(output_image_path, quality=95, format='JPEG')
|
||
else:
|
||
processed_img.save(output_image_path, quality=95, format='JPEG')
|
||
except Exception as jpeg_save_err:
|
||
logger.error(f"无法将 {output_filename} 保存为 JPEG: {jpeg_save_err}")
|
||
continue # 跳过这个图像
|
||
|
||
processed_count += 1
|
||
if (img_idx + 1) % 10 == 0 or (img_idx + 1) == len(selected_for_batch):
|
||
logger.info(f" 批次 {batch_index}: 已处理 {img_idx + 1}/{len(selected_for_batch)} 张图像...")
|
||
else:
|
||
logger.warning(f"处理图像失败: {original_filename}")
|
||
|
||
except FileNotFoundError:
|
||
logger.error(f"找不到输入图像文件: {input_image_path}")
|
||
except Image.UnidentifiedImageError:
|
||
logger.error(f"无法识别或打开图像文件: {input_image_path}")
|
||
except Exception as proc_err:
|
||
logger.error(f"处理图像时发生未知错误 '{original_filename}': {proc_err}")
|
||
|
||
logger.info(f"批次 {batch_index} 处理完成，成功处理 {processed_count} 张图像。保存在: {processed_output_dir}")
|
||
# --- 新增：处理选定的图像 --- END ---
|
||
else:
|
||
logger.warning(f"批次 {batch_index} 未选择任何图像。")
|
||
|
||
# 保存结果到JSON文件
|
||
try:
|
||
output_dir = os.path.dirname(args.output_file)
|
||
if output_dir and not os.path.exists(output_dir):
|
||
os.makedirs(output_dir)
|
||
logger.info(f"创建输出目录: {output_dir}")
|
||
|
||
with open(args.output_file, 'w', encoding='utf-8') as f:
|
||
json.dump(all_batches, f, indent=4, ensure_ascii=False)
|
||
logger.info(f"所有批次的选择结果已保存到: {args.output_file}")
|
||
except IOError as e:
|
||
logger.error(f"无法写入输出文件 {args.output_file}: {e}")
|
||
# 如果保存失败，可以选择在控制台打印结果
|
||
print("\n--- 选择结果 (保存失败，打印至控制台) ---") # 修复了字符串换行问题
|
||
print(json.dumps(all_batches, indent=4, ensure_ascii=False))
|
||
print("---------------------------------------")
|
||
except OSError as e:
|
||
logger.error(f"创建输出目录时出错: {e}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |