289 lines
8.9 KiB
Python
289 lines
8.9 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
生成图片模块
|
||
"""
|
||
|
||
# 列举导入模块
|
||
import cv2
|
||
import numpy
|
||
from pathlib import Path
|
||
import random
|
||
from typing import List, Dict, Optional, Union, Literal, Tuple
|
||
|
||
from numpy._core import int64
|
||
|
||
|
||
class Canvas:
|
||
"""
|
||
画布，用于绘制图片
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
materials_folder_path: str = r"E:\jianying\materials\纯图",
|
||
canvas_height: int = 1920,
|
||
canvas_width: int = 1080,
|
||
):
|
||
"""
|
||
初始化
|
||
:param materials_folder_path: 所有素材文件夹的根目录路径
|
||
:param canvas_height: 画布高度
|
||
:param canvas_width: 画布宽度
|
||
"""
|
||
# 构建所有素材文件夹的根目录路径
|
||
self.materials_folder_path = Path(materials_folder_path)
|
||
|
||
self.canvas_height, self.canvas_width = canvas_height, canvas_width
|
||
# 创建画布（默认尺寸为 1920x1080，底色为白色）
|
||
self.canvas = (
|
||
numpy.ones((self.canvas_height, self.canvas_width, 4), dtype=numpy.uint8)
|
||
* 255
|
||
)
|
||
|
||
def _get_materials_paths(self, folder_name: str) -> List[Path]:
|
||
"""
|
||
根据素材文件夹名称获取该文件夹内所有素材路径列表
|
||
:return: 该文件夹内所有素材路径列表
|
||
"""
|
||
# 构建指定素材文件路径
|
||
folder_path = self.materials_folder_path / folder_name
|
||
if not (folder_path.exists() and folder_path.is_dir()):
|
||
raise RuntimeError(f"{folder_name} 素材文件夹不存在")
|
||
# 获取该文件夹内所有素材路径列表
|
||
material_paths = list(folder_path.rglob(pattern="*.png"))
|
||
if not material_paths:
|
||
raise RuntimeError(f"{folder_name} 素材文件夹为空")
|
||
return material_paths
|
||
|
||
def _read_material(
|
||
self,
|
||
material_path: Path,
|
||
) -> numpy.ndarray:
|
||
"""
|
||
读取素材
|
||
:param material_path: 素素材文件路径
|
||
:return: 素材 OpenCV 数据
|
||
"""
|
||
with open(file=material_path, mode="rb") as file:
|
||
material_buf = numpy.frombuffer(
|
||
buffer=file.read(), dtype=numpy.uint8
|
||
) # 二进制数据
|
||
# 解码为图像 OpenCV 数据
|
||
material = cv2.imdecode(buf=material_buf, flags=cv2.IMREAD_UNCHANGED)
|
||
if material is None:
|
||
raise RuntimeError(f"读取素材失败")
|
||
return material
|
||
|
||
def _resize_material(
|
||
self,
|
||
material: numpy.ndarray,
|
||
scale: Union[Literal["fixed_width"], float] = 1.0,
|
||
) -> numpy.ndarray:
|
||
"""
|
||
缩放素材
|
||
:param material: 素材 OpenCV 数据
|
||
:param scale: 缩放比例
|
||
:return: 缩放后素材 OpenCV 数据
|
||
"""
|
||
# 获取素材高度和宽度
|
||
material_height, material_width = material.shape[:2]
|
||
|
||
match scale:
|
||
case "fixed_width":
|
||
material_height, material_width = (
|
||
int(material_height * self.canvas_width / material_width),
|
||
self.canvas_width,
|
||
)
|
||
case _:
|
||
material_height, material_width = int(material_height * scale), int(
|
||
material_width * scale
|
||
)
|
||
|
||
return cv2.resize(
|
||
src=material,
|
||
dsize=(material_width, material_height),
|
||
interpolation=cv2.INTER_LINEAR,
|
||
)
|
||
|
||
def _paste_material(
|
||
self,
|
||
material: numpy.ndarray,
|
||
position: Union[str, Tuple[Optional[int], Optional[int]]] = "center",
|
||
) -> None:
|
||
"""
|
||
粘贴素材
|
||
:param material: 素材 OpenCV 数据
|
||
:param position: 粘贴位置
|
||
:return: None
|
||
"""
|
||
# 获取素材高度和宽度
|
||
material_height, material_width = material.shape[:2]
|
||
|
||
# 计算素材在画布原点坐标
|
||
x = (self.canvas_width - material_width) // 2 # 水平方向坐标
|
||
y = (self.canvas_height - material_height) // 2 # 垂直方向坐标
|
||
|
||
match position:
|
||
case "center": # 居中
|
||
pass
|
||
case _:
|
||
if position[0] is not None:
|
||
x = position[0] - material_width // 2
|
||
if position[1] is not None:
|
||
y = position[1] - material_height // 2
|
||
|
||
# 计算素材左上角坐标
|
||
x1, y1 = max(x, 0), max(y, 0)
|
||
# 计算素材右下角坐标
|
||
x2, y2 = min(x + material_width, self.canvas_width), min(
|
||
y + material_height, self.canvas_height
|
||
)
|
||
# 截取素材有效区域
|
||
material = material[y1 - y : y2 - y, x1 - x : x2 - x]
|
||
|
||
# 取出画布对应区域
|
||
area = self.canvas[y1:y2, x1:x2]
|
||
|
||
# 透明度计算
|
||
alpha_fg = material[..., 3] / 255.0
|
||
alpha_bg = area[..., 3] / 255.0
|
||
|
||
# 颜色混合
|
||
for c in range(3):
|
||
area[..., c] = (
|
||
material[..., c] * alpha_fg + area[..., c] * alpha_bg * (1 - alpha_fg)
|
||
).astype(numpy.uint8)
|
||
|
||
# 新透明度
|
||
area[..., 3] = ((alpha_fg + alpha_bg * (1 - alpha_fg)) * 255).astype(
|
||
numpy.uint8
|
||
)
|
||
|
||
# 写回画布
|
||
self.canvas[y1:y2, x1:x2] = area
|
||
|
||
def _insert_text(
|
||
self,
|
||
text: str,
|
||
):
|
||
"""
|
||
插入文字
|
||
:param text: 文字内容
|
||
"""
|
||
# 字体
|
||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||
# 字体缩放比例
|
||
font_scale = 2
|
||
# 文字厚度
|
||
thickness = 3
|
||
|
||
# 获取文字尺寸
|
||
(text_width, text_height), _ = cv2.getTextSize(
|
||
text=text, fontFace=font, fontScale=font_scale, thickness=thickness
|
||
)
|
||
|
||
cv2.putText(
|
||
img=self.canvas,
|
||
text=text,
|
||
org=(
|
||
(self.canvas_width - text_width) // 2,
|
||
(self.canvas_height - text_height) // 2,
|
||
),
|
||
fontFace=font,
|
||
fontScale=font_scale,
|
||
color=(0, 0, 0, 255),
|
||
thickness=thickness,
|
||
lineType=cv2.LINE_AA,
|
||
)
|
||
|
||
def _save(
|
||
self,
|
||
image_path: str,
|
||
) -> None:
|
||
"""
|
||
保存图片
|
||
:param image_path: 图片保存路径
|
||
:return: None
|
||
"""
|
||
# 编码为二进制数据
|
||
success, image_encoded = cv2.imencode(
|
||
".jpg", img=self.canvas, params=[cv2.IMWRITE_JPEG_QUALITY, 95]
|
||
)
|
||
if success:
|
||
with open(file=(Path(__file__).parent / image_path), mode="wb") as file:
|
||
file.write(image_encoded.tobytes())
|
||
print("✅ 图片保存成功！")
|
||
else:
|
||
print("❌ 图片编码失败！")
|
||
|
||
def draw(self, folder_names: List[str], counts: int) -> None:
|
||
"""
|
||
绘制图片
|
||
:param folder_names: 绘制所用素材文件夹名称列表
|
||
:param counts: 绘制图片数量
|
||
:return: None
|
||
"""
|
||
materials_paths = {} # 绘制所用素材文件夹名称和其素材路径列表
|
||
for folder_name in folder_names:
|
||
materials_paths.update(
|
||
{folder_name: self._get_materials_paths(folder_name)}
|
||
)
|
||
|
||
current_count = 0 # 当前绘制图片数量计数
|
||
while current_count < counts:
|
||
self.canvas = (
|
||
numpy.ones(
|
||
(self.canvas_height, self.canvas_width, 4), dtype=numpy.uint8
|
||
)
|
||
* 255
|
||
)
|
||
|
||
# 竖纹素材
|
||
vertical_strip = self._read_material(
|
||
material_path=random.choice(materials_paths["竖纹"])
|
||
)
|
||
# 等宽缩放竖纹素材
|
||
vertical_strip = self._resize_material(
|
||
material=vertical_strip, scale="fixed_width"
|
||
)
|
||
# 粘贴竖纹素材到画布
|
||
self._paste_material(material=vertical_strip)
|
||
|
||
# 按钮素材
|
||
button = self._read_material(
|
||
material_path=random.choice(materials_paths["按钮"])
|
||
)
|
||
# 粘贴按钮素材到画布
|
||
self._paste_material(material=button)
|
||
|
||
# 插入文字
|
||
self._insert_text(text="还未登录")
|
||
|
||
# 箭头素材
|
||
arrow = self._read_material(
|
||
material_path=random.choice(materials_paths["箭头"])
|
||
)
|
||
# 粘贴箭头素材到画布
|
||
self._paste_material(material=arrow, position=(None, 1280))
|
||
|
||
# 声明素材
|
||
declaration = self._read_material(
|
||
material_path=random.choice(materials_paths["声明"])
|
||
)
|
||
# 粘贴声明素材到画布
|
||
self._paste_material(material=declaration, position=(None, 1640))
|
||
|
||
# 保存图片
|
||
self._save(
|
||
image_path=f"draw_{current_count}.jpg",
|
||
)
|
||
current_count += 1
|
||
|
||
|
||
canvas = Canvas()
|
||
|
||
canvas.draw(
|
||
folder_names=["竖纹", "按钮", "箭头", "声明"],
|
||
counts=1,
|
||
)
|