Create utils/bg_generator.py
Browse files- utils/bg_generator.py +85 -0
utils/bg_generator.py
ADDED
|
@@ -0,0 +1,85 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
from pathlib import Path
|
| 3 |
+
from typing import List, Tuple
|
| 4 |
+
import time, random
|
| 5 |
+
import numpy as np
|
| 6 |
+
from PIL import Image, ImageFilter, ImageOps
|
| 7 |
+
|
| 8 |
+
TMP_DIR = Path("/tmp/bgfx"); TMP_DIR.mkdir(parents=True, exist_ok=True)
|
| 9 |
+
|
| 10 |
+
_PALETTES = {
|
| 11 |
+
"office": [(240,245,250),(210,220,230),(180,190,200)],
|
| 12 |
+
"studio": [(18,18,20),(32,32,36),(58,60,64)],
|
| 13 |
+
"sunset": [(255,183,77),(255,138,101),(244,143,177)],
|
| 14 |
+
"forest": [(46,125,50),(102,187,106),(165,214,167)],
|
| 15 |
+
"ocean": [(33,150,243),(3,169,244),(0,188,212)],
|
| 16 |
+
"minimal": [(245,246,248),(230,232,236),(214,218,224)],
|
| 17 |
+
"warm": [(255,224,178),(255,204,128),(255,171,145)],
|
| 18 |
+
"cool": [(197,202,233),(179,229,252),(178,235,242)],
|
| 19 |
+
"royal": [(63,81,181),(121,134,203),(159,168,218)],
|
| 20 |
+
}
|
| 21 |
+
|
| 22 |
+
def _save_pil(img: Image.Image, stem: str = "ai_bg", ext: str = "png") -> str:
|
| 23 |
+
ts = int(time.time() * 1000)
|
| 24 |
+
p = TMP_DIR / f"{stem}_{ts}.{ext}"
|
| 25 |
+
img.save(p)
|
| 26 |
+
return str(p)
|
| 27 |
+
|
| 28 |
+
def _palette_from_prompt(prompt: str) -> List[tuple]:
|
| 29 |
+
p = (prompt or "").lower()
|
| 30 |
+
for key, pal in _PALETTES.items():
|
| 31 |
+
if key in p:
|
| 32 |
+
return pal
|
| 33 |
+
random.seed(hash(p) % (2**32 - 1))
|
| 34 |
+
return [tuple(random.randint(90, 200) for _ in range(3)) for _ in range(3)]
|
| 35 |
+
|
| 36 |
+
def _perlin_like_noise(h: int, w: int, octaves: int = 4) -> np.ndarray:
|
| 37 |
+
acc = np.zeros((h, w), dtype=np.float32)
|
| 38 |
+
for o in range(octaves):
|
| 39 |
+
scale = 2 ** o
|
| 40 |
+
small = np.random.rand(h // scale + 1, w // scale + 1).astype(np.float32)
|
| 41 |
+
small = Image.fromarray((small * 255).astype(np.uint8)).resize((w, h), Image.BILINEAR)
|
| 42 |
+
acc += np.array(small, dtype=np.float32) / 255.0 / (o + 1)
|
| 43 |
+
acc /= max(1e-6, acc.max())
|
| 44 |
+
return acc
|
| 45 |
+
|
| 46 |
+
def _blend_palette(noise: np.ndarray, palette: List[tuple]) -> Image.Image:
|
| 47 |
+
h, w = noise.shape
|
| 48 |
+
img = np.zeros((h, w, 3), dtype=np.float32)
|
| 49 |
+
t1, t2 = 0.33, 0.66
|
| 50 |
+
c0, c1, c2 = [np.array(c, dtype=np.float32) for c in palette]
|
| 51 |
+
m0, m1, m2 = noise < t1, (noise >= t1) & (noise < t2), noise >= t2
|
| 52 |
+
img[m0], img[m1], img[m2] = c0, c1, c2
|
| 53 |
+
return Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))
|
| 54 |
+
|
| 55 |
+
def generate_ai_background(
|
| 56 |
+
prompt: str, width: int = 1280, height: int = 720,
|
| 57 |
+
bokeh: float = 0.0, vignette: float = 0.15, contrast: float = 1.05
|
| 58 |
+
) -> Tuple[Image.Image, str]:
|
| 59 |
+
palette = _palette_from_prompt(prompt)
|
| 60 |
+
noise = _perlin_like_noise(height, width, octaves=4)
|
| 61 |
+
img = _blend_palette(noise, palette)
|
| 62 |
+
|
| 63 |
+
if bokeh > 0:
|
| 64 |
+
img = img.filter(ImageFilter.GaussianBlur(radius=max(0, min(50, bokeh))))
|
| 65 |
+
|
| 66 |
+
if vignette > 0:
|
| 67 |
+
import numpy as np
|
| 68 |
+
base = np.array(img).astype(np.float32) / 255.0
|
| 69 |
+
y, x = np.ogrid[:height, :width]
|
| 70 |
+
cx, cy = width / 2, height / 2
|
| 71 |
+
r = np.sqrt((x - cx) ** 2 + (y - cy) ** 2)
|
| 72 |
+
mask = 1 - np.clip(r / (max(width, height) / 1.2), 0, 1)
|
| 73 |
+
mask = (mask ** 2) * (1 - vignette) + vignette
|
| 74 |
+
out = base * mask[..., None]
|
| 75 |
+
img = Image.fromarray(np.clip(out * 255, 0, 255).astype(np.uint8))
|
| 76 |
+
|
| 77 |
+
if contrast != 1.0:
|
| 78 |
+
img = ImageOps.autocontrast(img, cutoff=1)
|
| 79 |
+
arr = np.array(img).astype(np.float32)
|
| 80 |
+
mean = arr.mean(axis=(0, 1), keepdims=True)
|
| 81 |
+
arr = (arr - mean) * float(contrast) + mean
|
| 82 |
+
img = Image.fromarray(np.clip(arr, 0, 255).astype(np.uint8))
|
| 83 |
+
|
| 84 |
+
path = _save_pil(img)
|
| 85 |
+
return img, path
|