添加卵总展示

konabot/plugins/memepack/__init__.py

@@ -2,11 +2,11 @@ from io import BytesIO
 from typing import Iterable, cast
 
 from nonebot import on_message
-from nonebot_plugin_alconna import (Alconna, Args, Field, MultiVar, Text,
+from nonebot_plugin_alconna import (Alconna, Args, Field, Image, MultiVar, Text,
                                     UniMessage, UniMsg, on_alconna)
 
-from konabot.common.nb.extract_image import extract_image_from_message
-from konabot.plugins.memepack.drawing.display import draw_cao_display
+from konabot.common.nb.extract_image import PIL_Image, extract_image_from_message
+from konabot.plugins.memepack.drawing.display import draw_cao_display, draw_snaur_display
 from konabot.plugins.memepack.drawing.saying import (draw_cute_ten,
                                                      draw_geimao, draw_mnk,
                                                      draw_pt, draw_suan)
@@ -139,3 +139,17 @@ async def _(msg: UniMsg, evt: Event, bot: Bot):
                     .text(err)
                     .export()
             )
+
+
+snaur_display_cmd = on_alconna(Alconna(
+    "卵总展示",
+    Args["image", Image | None],
+))
+
+@snaur_display_cmd.handle()
+async def _(img: PIL_Image):
+    img_processed = await draw_snaur_display(img)
+    img_data = BytesIO()
+    img_processed.save(img_data, "PNG")
+    await snaur_display_cmd.send(await UniMessage().image(raw=img_data).export())
+

konabot/plugins/memepack/drawing/display.py

@@ -4,10 +4,12 @@ from typing import Any, cast
 import cv2
 import numpy as np
 import PIL.Image
+import PIL.ImageChops
+import PIL.ImageEnhance
 
 from konabot.common.path import ASSETS_PATH
 
-cao_image = PIL.Image.open(ASSETS_PATH / "img" / "meme" / "caoimg1.png")
+cao_image       = PIL.Image.open(ASSETS_PATH / "img" / "meme" / "caoimg1.png")
 CAO_QUAD_POINTS = np.float32(cast(Any, [
     [392, 540],
     [577, 557],
@@ -15,6 +17,16 @@ CAO_QUAD_POINTS = np.float32(cast(Any, [
     [381, 687],
 ]))
 
+snaur_image_base  = PIL.Image.open(ASSETS_PATH / "img" / "meme" / "snaur_1_base.png")
+snaur_image_top   = PIL.Image.open(ASSETS_PATH / "img" / "meme" / "snaur_1_top.png")
+SNAUR_RATIO       = (1 / 2) ** .5
+SNAUR_QUAD_POINTS = np.float32(cast(Any, [
+    [0,   466 ],
+    [673, 471 ],
+    [640, 1196],
+    [106, 1280],
+]))
+
 def _draw_cao_display(image: PIL.Image.Image):
     src = np.array(image.convert("RGB"))
     h, w = src.shape[:2]
@@ -43,3 +55,63 @@ def _draw_cao_display(image: PIL.Image.Image):
 
 async def draw_cao_display(image: PIL.Image.Image):
     return await asyncio.to_thread(_draw_cao_display, image)
+
+
+def _draw_snaur_display(image: PIL.Image.Image):
+    src = np.array(image.convert("RGB"))
+    _h, _w = src.shape[:2]
+
+    if _w / _h < SNAUR_RATIO:
+        _w_target = _w
+        _h_target = int(_w / SNAUR_RATIO)
+    else:
+        _w_target = int(_h * SNAUR_RATIO)
+        _h_target = _h
+
+    x_center = _w / 2
+    y_center = _h / 2
+
+    x1 = int(x_center - _w_target / 2)
+    x2 = int(x_center + _w_target / 2)
+    y1 = int(y_center - _h_target / 2)
+    y2 = int(y_center + _h_target / 2)
+
+    src = src[y1:y2, x1:x2, :]
+
+    h, w = src.shape[:2]
+    src_points = np.float32(cast(Any, [
+        [0, 0],
+        [w, 0],
+        [w, h],
+        [0, h],
+    ]))
+    dst_points = SNAUR_QUAD_POINTS
+    M = cv2.getPerspectiveTransform(cast(Any, src_points), cast(Any, dst_points))
+    output_size = snaur_image_top.size
+    output_w, output_h = output_size
+    warped = cv2.warpPerspective(
+        src,
+        M,
+        (output_w, output_h),
+        flags=cv2.INTER_LINEAR,
+        borderMode=cv2.BORDER_CONSTANT,
+        borderValue=(0, 0, 0)
+    )
+
+    result = PIL.Image.fromarray(warped, 'RGB').convert('RGBA')
+    result = PIL.ImageChops.multiply(result, snaur_image_base)
+
+    r, g, b, a = result.split()
+    a = a.point(lambda p: int(p * 0.8))
+    result = PIL.Image.merge('RGBA', (r, g, b, a))
+
+    enhancer = PIL.ImageEnhance.Color(result)
+    result = enhancer.enhance(0.85)
+
+    result = PIL.Image.alpha_composite(result, snaur_image_top)
+    return result
+
+
+async def draw_snaur_display(image: PIL.Image.Image) -> PIL.Image.Image:
+    return await asyncio.to_thread(_draw_snaur_display, image)
+