gacha-bot2 / utils /generators /custom.py
Amit9011's picture
Update utils/generators/custom.py
78a11e2 verified
Raw
History Blame Contribute Delete
5.99 kB
# utils/generators/custom.py
import io
import logging
from PIL import Image, ImageDraw, ImageOps, ImageFilter, ImageStat
from .base import BaseGenerator
logger = logging.getLogger(__name__)
class CustomGenerator(BaseGenerator):
def build_template(self, image_data: bytes, frame_data: bytes, char_name: str) -> Image.Image:
cleanup = []
try:
canvas = Image.new("RGBA", (self.CANVAS_W, self.CANVAS_H), (0, 0, 0, 0))
# 1. Load the frame and IMMEDIATELY resize it to the exact card bounds (580x820)
# This guarantees we never double-resample the character later.
frame_raw = Image.open(io.BytesIO(frame_data)).convert("RGBA")
cleanup.append(frame_raw)
frame_base = frame_raw.resize((self.CARD_W, self.CARD_H), Image.Resampling.LANCZOS)
cleanup.append(frame_base)
f_w, f_h = self.CARD_W, self.CARD_H
alpha = frame_base.split()[3]
# 2. Extract dynamic theme color
try:
solid_pixels = alpha.point(lambda p: 255 if p > 50 else 0, mode="1")
stat = ImageStat.Stat(frame_base.convert("RGB"), mask=solid_pixels)
avg_r, avg_g, avg_b = [int(x) for x in stat.mean]
if max(avg_r, avg_g, avg_b) < 80:
avg_r, avg_g, avg_b = 160, 160, 160
self.frame_color = (avg_r, avg_g, avg_b)
except Exception as e:
logger.error(f"Color extraction failed: {e}")
self.frame_color = None
# 3. Create a STRICT binary mask for the floodfill
# This prevents the "diagonal slice" by removing soft, semi-transparent pixels
binary_alpha = alpha.point(lambda p: 255 if p > 20 else 0, mode="L")
cleanup.append(binary_alpha)
ImageDraw.floodfill(binary_alpha, (f_w // 2, f_h // 2), 128)
hole_mask = binary_alpha.point(lambda p: 255 if p == 128 else 0, mode="L")
cleanup.append(hole_mask)
# Anti-alias the final cutout so edges are smooth
bleed_mask = hole_mask.filter(ImageFilter.MaxFilter(5)).filter(ImageFilter.GaussianBlur(1.5))
cleanup.append(bleed_mask)
# 4. Load Character
char_raw = Image.open(io.BytesIO(image_data)).convert("RGBA")
cleanup.append(char_raw)
bbox = hole_mask.getbbox()
if bbox:
hx1, hy1, hx2, hy2 = bbox
hw, hh = hx2 - hx1, hy2 - hy1
else:
hx1, hy1, hw, hh = 0, 0, f_w, f_h
# 5. Crop character EXACTLY ONCE to the final 1:1 pixel dimensions (LANCZOS for max quality)
char_fitted = ImageOps.fit(char_raw, (hw, hh), method=Image.Resampling.LANCZOS, centering=(0.5, 0.25))
cleanup.append(char_fitted)
# 6. Composite
char_full = Image.new("RGBA", (f_w, f_h), (0, 0, 0, 0))
char_full.paste(char_fitted, (hx1, hy1))
cleanup.append(char_full)
char_full.putalpha(bleed_mask)
composite = Image.new("RGBA", (f_w, f_h), (0, 0, 0, 0))
composite.paste(char_full, (0, 0), char_full)
composite.paste(frame_base, (0, 0), frame_base)
cleanup.append(composite)
# 7. Final Paste (Card fits perfectly on the 660x900 canvas with no extra resizing)
canvas.paste(composite, (self.OFFSET_X, self.OFFSET_Y), composite)
return canvas
except Exception as e:
logger.error(f"Cookie Cutter Frame Error: {e}", exc_info=True)
return Image.new("RGBA", (self.CANVAS_W, self.CANVAS_H), (0,0,0,0))
finally:
for img in cleanup:
try: img.close()
except: pass
def apply_text(self, canvas: Image.Image, char_name: str, print_num: int, element: str, rarity: str) -> io.BytesIO:
uid = ImageDraw.Draw(canvas, "RGBA")
card_left = self.OFFSET_X
card_right = self.OFFSET_X + self.CARD_W
card_bottom = self.OFFSET_Y + self.CARD_H
card_w = self.CARD_W
name_sz = max(24, int(card_w * 0.080))
min_sz = 20
name_font = self.get_font(name_sz)
num_font = self.get_font(int(name_sz * 0.95))
r_color = self.ELEMENT_COLORS.get(element, (168, 85, 247))
glow_solid = getattr(self, 'frame_color', r_color[:3])
final_char_name = char_name.title()
# Keep text strictly inside the borders
max_name_w = int(card_w * 0.85)
# Dynamic Text Scaler Loop
while True:
try:
n_bbox = uid.textbbox((0, 0), final_char_name, font=name_font)
except:
n_bbox = (0, 0, 20, name_sz)
n_w = n_bbox[2] - n_bbox[0]
if n_w <= max_name_w or name_sz <= min_sz:
break
name_sz -= 2
name_font = self.get_font(name_sz)
name_x = card_left + (card_w - n_w) // 2
name_y = card_bottom - 120
num_text = f"{print_num}"
try: num_bbox = uid.textbbox((0, 0), num_text, font=num_font)
except: num_bbox = (0, 0, 40, 30)
num_w = num_bbox[2] - num_bbox[0]
num_h = num_bbox[3] - num_bbox[1]
right_margin = max(40, int(card_w * 0.18))
num_x = card_right - right_margin - num_w
num_x = max(card_left + right_margin, num_x)
num_y = name_y - num_h - 10
self.draw_text_with_glow(uid, (name_x, name_y), final_char_name, name_font, core_color="white", glow_color=glow_solid)
self.draw_text_with_glow(uid, (num_x, num_y), num_text, num_font, core_color="white", glow_color=glow_solid)
buf = io.BytesIO()
canvas.save(buf, format="WEBP", quality=95)
buf.seek(0)
return buf