model_train_dm/backend/apps/common/utils/watermarker.py

import enum
import math
import os
import textwrap
from typing import Union

from PIL import Image, ImageChops, ImageDraw, ImageEnhance, ImageFont, ImageOps


class WatermarkerStyles(enum.Enum):
    """水印样式"""

    STRIPED = 1
    CENTRAL = 2


class Watermarker(object):
    """图片水印工具"""

    django_support = False

    def __init__(
        self,
        image: Union[str, Image.Image],
        text: str,
        style: WatermarkerStyles.CENTRAL,
        angle=25,
        color="black",
        font_file="汉仪旗黑.ttf",
        font_height_crop=1.2,
        opacity=0.15,
        quality=100,
        size=50,
        space=100,
        chars_per_line=10,
    ):
        if isinstance(image, str):
            self.image_path = image
            self.origin_image = Image.open(image)
        elif isinstance(image, Image.Image):
            self.origin_image = image
            self.image_path = None
        else:
            raise ValueError("`image` 参数必须是图片路径或 PIL Image 对象")
        self.text = text
        self.style = style
        self.angle = angle
        self.color = color

        if self.django_support:
            from django.conf import settings

            self.font_file = os.path.join(settings.BASE_DIR, "contrib", "image", "font", font_file)
        else:
            self.font_file = os.path.join(os.path.abspath("."), "font", font_file)

        self.font_height_crop = font_height_crop
        self.opacity = opacity
        self.quality = quality
        self.size = size
        self.space = space
        self._result_image = None
        self.chars_per_line = chars_per_line

    @staticmethod
    def set_image_opacity(image: Image, opacity: float):
        assert 0 <= opacity <= 1

        alpha = image.split()[3]
        alpha = ImageEnhance.Brightness(alpha).enhance(opacity)
        image.putalpha(alpha)
        return image

    @staticmethod
    def crop_image_edge(image: Image):
        bg = Image.new(mode="RGBA", size=image.size)
        diff = ImageChops.difference(image, bg)
        del bg
        bbox = diff.getbbox()
        if bbox:
            return image.crop(bbox)
        return image

    def _add_mark_striped(self):
        origin_image = self.origin_image
        origin_image = ImageOps.exif_transpose(origin_image)

        width = len(self.text) * self.size
        height = round(self.size * self.font_height_crop)

        watermark_image = Image.new(mode="RGBA", size=(width, height))

        draw_table = ImageDraw.Draw(im=watermark_image)
        draw_table.text(
            xy=(0, 0),
            text=self.text,
            fill=self.color,
            font=ImageFont.truetype(self.font_file, size=self.size),
        )
        del draw_table

        watermark_image = Watermarker.crop_image_edge(watermark_image)

        Watermarker.set_image_opacity(watermark_image, self.opacity)

        c = int(
            math.sqrt(
                origin_image.size[0] * origin_image.size[0] + origin_image.size[1] * origin_image.size[1]
            )
        )

        watermark_mask = Image.new(mode="RGBA", size=(c, c))

        y, idx = 0, 0
        while y < c:
            x = -int((watermark_image.size[0] + self.space) * 0.5 * idx)
            idx = (idx + 1) % 2

            while x < c:
                watermark_mask.paste(watermark_image, (x, y))
                x = x + watermark_image.size[0] + self.space
            y = y + watermark_image.size[1] + self.space

        watermark_mask = watermark_mask.rotate(self.angle)

        if origin_image.mode != "RGBA":
            origin_image = origin_image.convert("RGBA")
        origin_image.paste(
            watermark_mask,
            (int((origin_image.size[0] - c) / 2), int((origin_image.size[1] - c) / 2)),
            mask=watermark_mask.split()[3],
        )
        del watermark_mask

        return origin_image

    def _add_mark_central(self):
        origin_image = self.origin_image
        origin_image = ImageOps.exif_transpose(origin_image)

        text_lines = textwrap.wrap(self.text, width=self.chars_per_line)
        text = "\n".join(text_lines)

        width = len(text) * self.size
        height = round(self.size * self.font_height_crop * len(text_lines))

        watermark_image = Image.new(mode="RGBA", size=(width, height))

        draw_table = ImageDraw.Draw(im=watermark_image)
        draw_table.text(
            xy=(0, 0),
            text=text,
            fill=self.color,
            font=ImageFont.truetype(self.font_file, size=self.size),
        )
        del draw_table

        watermark_image = Watermarker.crop_image_edge(watermark_image)

        Watermarker.set_image_opacity(watermark_image, self.opacity)

        c = int(
            math.sqrt(
                origin_image.size[0] * origin_image.size[0] + origin_image.size[1] * origin_image.size[1]
            )
        )
        watermark_mask = Image.new(mode="RGBA", size=(c, c))
        watermark_mask.paste(
            watermark_image,
            (
                int((watermark_mask.width - watermark_image.width) / 2),
                int((watermark_mask.height - watermark_image.height) / 2),
            ),
        )
        watermark_mask = watermark_mask.rotate(self.angle)

        if origin_image.mode != "RGBA":
            origin_image = origin_image.convert("RGBA")

        box = (int((origin_image.width - watermark_mask.width) / 2), int((origin_image.height - watermark_mask.height) / 2))
        origin_image.paste(watermark_mask, box, mask=watermark_mask.split()[3])

        return origin_image

    @property
    def image(self):
        if not self._result_image:
            if self.style == WatermarkerStyles.STRIPED:
                self._result_image = self._add_mark_striped()
            if self.style == WatermarkerStyles.CENTRAL:
                self._result_image = self._add_mark_central()
        return self._result_image

    def save(self, file_path: str, image_format: str = "png"):
        with open(file_path, "wb") as f:
            self.image.save(f, image_format)

    def show(self):
        self.image.show()


if __name__ == "__main__":
    Watermarker("path_to_image", "mark text", WatermarkerStyles.CENTRAL).show()
    Watermarker("path_to_image", "mark text", WatermarkerStyles.CENTRAL).save("path_to_save")
    image = Watermarker("path_to_image", "mark text", WatermarkerStyles.CENTRAL).image