MogensR commited on
Commit
ef24f03
·
1 Parent(s): 66fefac

Create utils/bg_generator.py

Browse files
Files changed (1) hide show
  1. 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