26408

短视频合成自动化/image.py

@@ -7,20 +7,282 @@
 import cv2
 import numpy
 from pathlib import Path
+import random
+from typing import List, Dict, Optional, Union, Literal, Tuple
 
-# 创建画布（默认填充为白色）
+from numpy._core import int64
-canvas = numpy.ones((1080, 1920, 3), dtype=numpy.uint8) * 255
 
-# 当前文件夹路径
-current_path = Path(__file__).parent
 
-# 构建生成图片路径
+class Canvas:
-image_path = current_path / "111.jpg"
+    """
+    画布，用于绘制图片
+    """
 
-print(image_path)
+    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)
 
-success = cv2.imwrite(
+        self.canvas_height, self.canvas_width = canvas_height, canvas_width
-    filename=image_path, img=canvas, params=[cv2.IMWRITE_JPEG_QUALITY, 95]
+        # 创建画布（默认尺寸为 1920x1080，底色为白色）
+        self.canvas = (
+            numpy.ones((self.canvas_height, self.canvas_width, 4), dtype=numpy.uint8)
+            * 255
         )
 
-print(success)
+    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,
+)