132 lines
5.2 KiB
Python
132 lines
5.2 KiB
Python
import os
|
||
import random
|
||
import shutil
|
||
from PIL import Image, ImageEnhance
|
||
|
||
|
||
class ImageExtendProcessor:
|
||
def is_image_file(self, filename):
|
||
"""检查文件是否为图像文件"""
|
||
image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff')
|
||
return filename.lower().endswith(image_extensions)
|
||
|
||
def random_cute_image(self, image_path, output_path, ratio=0.8):
|
||
"""
|
||
对图像进行随机裁剪
|
||
:param image_path: 输入图像路径
|
||
:param output_path: 输出图像路径
|
||
:param ratio: 裁剪比例，决定裁剪区域的大小，默认为0.8
|
||
"""
|
||
try:
|
||
with Image.open(image_path) as img:
|
||
# 获取图像尺寸
|
||
width, height = img.size
|
||
|
||
# 计算裁剪后的尺寸
|
||
new_width = int(width * ratio)
|
||
new_height = int(height * ratio)
|
||
|
||
# 随机生成裁剪起始点
|
||
left = random.randint(0, width - new_width)
|
||
top = random.randint(0, height - new_height)
|
||
right = left + new_width
|
||
bottom = top + new_height
|
||
|
||
# 执行裁剪
|
||
cropped_img = img.crop((left, top, right, bottom))
|
||
|
||
# 保存裁剪后的图像
|
||
cropped_img.save(output_path)
|
||
return True
|
||
except Exception as e:
|
||
print(f"处理图像 {image_path} 时出错: {e}")
|
||
return False
|
||
|
||
def random_brightness(self, image_path, output_path, brightness_factor=None):
|
||
"""
|
||
对图像进行随机亮度调整
|
||
:param image_path: 输入图像路径
|
||
:param output_path: 输出图像路径
|
||
:param brightness_factor: 亮度调整因子，默认为随机值
|
||
"""
|
||
try:
|
||
with Image.open(image_path) as img:
|
||
# 创建一个ImageEnhance.Brightness对象
|
||
enhancer = ImageEnhance.Brightness(img)
|
||
|
||
# 如果没有指定亮度因子，则随机生成
|
||
if brightness_factor is None:
|
||
brightness_factor = random.uniform(0.5, 1.5)
|
||
|
||
# 应用亮度调整
|
||
brightened_img = enhancer.enhance(brightness_factor)
|
||
|
||
# 保存调整后的图像
|
||
brightened_img.save(output_path)
|
||
return True
|
||
except Exception as e:
|
||
print(f"处理图像 {image_path} 时出错: {e}")
|
||
return False
|
||
|
||
def rotate_image(self, image_path, output_path, degrees):
|
||
"""旋转图像并保存到指定路径"""
|
||
try:
|
||
with Image.open(image_path) as img:
|
||
# 旋转图像并自动调整画布大小
|
||
rotated = img.rotate(degrees, expand=True)
|
||
rotated.save(output_path)
|
||
return True
|
||
except Exception as e:
|
||
print(f"处理图像 {image_path} 时出错: {e}")
|
||
return False
|
||
|
||
def process_image_directory(self, src_dir, dst_dir, same_directory, **kwargs):
|
||
"""
|
||
处理单个目录中的图像文件
|
||
:param src_dir: 源目录路径
|
||
:param dst_dir: 目标目录路径
|
||
"""
|
||
if not os.path.exists(dst_dir):
|
||
os.makedirs(dst_dir)
|
||
# 获取目录中所有图像文件
|
||
image_files = [f for f in os.listdir(src_dir)
|
||
if self.is_image_file(f) and os.path.isfile(os.path.join(src_dir, f))]
|
||
|
||
# 处理每个图像文件
|
||
for img_file in image_files:
|
||
src_path = os.path.join(src_dir, img_file)
|
||
base_name, ext = os.path.splitext(img_file)
|
||
if not same_directory:
|
||
# 复制原始文件 （另存文件夹时启用）
|
||
shutil.copy2(src_path, os.path.join(dst_dir, img_file))
|
||
|
||
# 生成并保存旋转后的图像
|
||
for angle in [90, 180, 270]:
|
||
dst_path = os.path.join(dst_dir, f"{base_name}_rotated_{angle}{ext}")
|
||
self.rotate_image(src_path, dst_path, angle)
|
||
for ratio in [0.8, 0.85, 0.9]:
|
||
dst_path = os.path.join(dst_dir, f"{base_name}_cute_{ratio}{ext}")
|
||
self.random_cute_image(src_path, dst_path, ratio)
|
||
for brightness_factor in [0.8, 0.9, 1.0]:
|
||
dst_path = os.path.join(dst_dir, f"{base_name}_brightness_{brightness_factor}{ext}")
|
||
self.random_brightness(src_path, dst_path, brightness_factor)
|
||
|
||
def image_extend(self, src_dir, dst_dir, same_directory=False, **kwargs):
|
||
if same_directory:
|
||
n_dst_dir = src_dir
|
||
print(f"处理目录 {src_dir} 中的图像文件 保存至同一目录下")
|
||
else:
|
||
n_dst_dir = dst_dir
|
||
print(f"处理目录 {src_dir} 中的图像文件 保存至不同目录下")
|
||
for src_subdir in os.listdir(src_dir):
|
||
src_subdir_path = os.path.join(src_dir, src_subdir)
|
||
dst_subdir_path = os.path.join(n_dst_dir, src_subdir)
|
||
self.process_image_directory(src_subdir_path, dst_subdir_path, same_directory)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
src_dir = "./scatter_mini"
|
||
dst_dir = "./scatter_add"
|
||
image_ex = ImageExtendProcessor()
|
||
image_ex.image_extend(src_dir, dst_dir, same_directory=False)
|