This seems like a fun one, and it might yield some interesting results.
I'll run masked training through sd-scripts with auto-generated masks for regularization.
I figure if I regenerate them every epoch using random seeds it should produce some interesting outcomes.
Place your image_name.png and image_name.txt in the folder side-by-side and give her a try.
The colab/jupyter script is here. Should just run from terminal and install using the pips. Install the correct version of cuda for your device if you don't have it installed already in your environment.
If you're on gpu install pytorch GPU - should be cuda121 on runpod.
https://gist.github.com/Birch-san/211f31f8d901dadd1025398fa1a603b8To install deepghs for gpu use;
!pip install dghs-imgutils #[gpu]
!pip install pillow scipy numpy tqdm transformers torch#!/usr/bin/env python3
"""
spiral_depth_generator.py • v18
─────────────────────────────────────────────────────────────
• 400 spirals per layer, sprites scaled bigger via spiral_scale
• Hands → smooth enlarged dip, Faces → smooth bump
• Depth model: Intel/dpt-hybrid-midas
----------------------------------------------------------------
pip install pillow scipy numpy tqdm transformers torch
"""
import math, random, gc
from pathlib import Path
from dataclasses import dataclass
from typing import Tuple, Optional
import numpy as np
from PIL import Image, ImageDraw
from scipy.ndimage import gaussian_filter
try:
from tqdm import tqdm
except ImportError:
tqdm = None
from imgutils.detect import detect_hands, detect_faces
# ───────── CONFIG ─────────────────────────────────────────
@dataclass
class Config:
images_dir: str = "./images"
output_dir: str = "spiral_output"
grid_size: int = 25
num_samples: Optional[int] = None
layers: Tuple[str, ...] = (
"ground","background1","background2",
"immediate","foreground1","foreground2","overlay")
layer_colors: Optional[dict] = None
spiral_type: str = "log"
tightness_range: Tuple[float,float] = (0.02, 0.14)
turns_range: Tuple[float,float] = (1.5, 15.5)
spirals_per_layer: int = 80
spiral_scale: float = 1.8 # ← new: sprite size vs. cell
blur_sigma: float = 1.0
alpha_gamma: float = 0.9
background_weight: float = 0.15
CLEANUP_GENERATED: bool = True
hand_mask_prob: float = 1.0
hand_floor: float = 0.3
hand_expand: float = 2.5
face_peak: float = 1.5
depth_ckpt: str = "Intel/dpt-hybrid-midas"
DEFAULT_COLORS = {
"ground": ( 80, 60, 60,255),
"background1": (100, 100, 150,180),
"background2": ( 90, 130, 110,160),
"immediate": (200, 180, 100,200),
"foreground1": (250, 150, 50,255),
"foreground2": (220, 220, 70,240),
"overlay": (255, 255, 255,180),
}
# ───────── DEPTH MODEL (unchanged) ────────────────────────
from transformers import AutoImageProcessor, DPTForDepthEstimation
import torch
class DepthModel:
def __init__(self, repo:str):
self.proc = AutoImageProcessor.from_pretrained(repo)
self.mdl = DPTForDepthEstimation.from_pretrained(repo).eval()
def infer(self, img):
t=self.proc(images=img, return_tensors="pt")
with torch.no_grad():
d=self.mdl(**t).predicted_depth.squeeze().cpu().numpy()
rng=d.max()-d.min()
d=(d-d.min())/rng if rng>=1e-6 else d
d=torch.nn.functional.interpolate(
torch.from_numpy(d[None,None]),
size=img.size[::-1], mode="bicubic", align_corners=False
)[0,0].numpy()
return np.clip(np.nan_to_num(d),0,1)
def unload(self):
del self.proc, self.mdl
gc.collect()
if torch.cuda.is_available(): torch.cuda.empty_cache()
# ───────── MISC HELPERS ───────────────────────────────────
def draw_spiral(w,h,c,stype,tight,turns):
img=Image.new("RGBA",(w,h),(0,0,0,0)); d=ImageDraw.Draw(img)
cx,cy=w//2,h//2; maxr=min(cx,cy); theta=0.0; prev=None
while theta<=turns*2*math.pi:
r=(math.exp(tight*theta) if stype=="log" else 5+4*theta if stype=="arch"
else math.sqrt(theta)*10 if stype=="fermat"
else 5*math.exp(math.log((1+math.sqrt(5))/2)/math.pi*theta))
if r>maxr: break
x,y=cx+r*math.cos(theta), cy+r*math.sin(theta)
a=int(c[3]*max(0,1-(r/maxr)**2))
if prev and a: d.line([prev,(x,y)], fill=(*c[:3],a), width=2)
prev=(x,y); theta+=0.05
return img
def dim_base(im,alpha=0.3):
r,g,b,a=im.convert("RGBA").split()
return Image.merge("RGBA",(r,g,b, Image.eval(a,lambda p:int(p*alpha))))
def gaussian_bump(h,w):
yy,xx=np.mgrid[0:h,0:w]
cy,cx=(h-1)/2,(w-1)/2
sigma=0.35*min(h,w)
return np.exp(-((yy-cy)**2+(xx-cx)**2)/(2*sigma*sigma)).astype(np.float32)
# ───────── GENERATOR ──────────────────────────────────────
class SpiralDepthGen:
def __init__(self,cfg:Config):
self.cfg=cfg
if not cfg.layer_colors: cfg.layer_colors=DEFAULT_COLORS
Path(cfg.output_dir).mkdir(parents=True, exist_ok=True)
self.depther=DepthModel(cfg.depth_ckpt)
self.bands=np.linspace(0,1,len(cfg.layers)+1)
self.seq=1
@staticmethod
def _hands(img): return [h[0] for h in detect_hands(img) or []]
@staticmethod
def _faces(img): return [f[0] for f in detect_faces(img) or []]
def _grid_depth(self,d,G):
H,W=d.shape; ch,cw=H/G,W/G
return np.array([[d[int((i+.5)*ch),int((j+.5)*cw)] for j in range(G)] for i in range(G)])
def _jitter(self, px, py, cw, ch):
return (int(px + random.uniform(-0.5,0.5)*cw),
int(py + random.uniform(-0.5,0.5)*ch))
def _process(self,p:Path):
base=Image.open(p).convert("RGB"); W,H=base.size
depth=self.depther.infer(base)
G=self.cfg.grid_size; cw,ch=W//G,H//G
gdepth=self._grid_depth(depth,G)
mask=np.power(depth,self.cfg.alpha_gamma)*self.cfg.background_weight*255
overlay=Image.new("RGBA",(W,H),(0,0,0,0)); caps=[]
for li,layer in enumerate(self.cfg.layers):
low,high=self.bands[li],self.bands[li+1]
cells=[(i,j) for i in range(G) for j in range(G) if low<=gdepth[i,j]<high]
random.shuffle(cells)
occ=np.zeros((G,G),bool); placed=0
for i,j in cells:
if placed>=self.cfg.spirals_per_layer: break
if occ[i,j]: continue
occ[i,j]=True
px,py=j*cw,i*ch
px,py=self._jitter(px,py,cw,ch)
# Spiral sprite scaled larger than cell
tile_w=int(cw*self.cfg.spiral_scale)
tile_h=int(ch*self.cfg.spiral_scale)
tile=draw_spiral(
tile_w,tile_h,
self.cfg.layer_colors[layer],
self.cfg.spiral_type,
random.uniform(*self.cfg.tightness_range),
random.uniform(*self.cfg.turns_range))
# top-left paste coordinate (centre on jitter point)
x0=px-tile_w//2; y0=py-tile_h//2
# clamp inside canvas
x0=max(0,min(x0,W-tile_w)); y0=max(0,min(y0,H-tile_h))
overlay.paste(tile,(x0,y0),tile)
alpha=np.array(tile.split()[-1],dtype=np.float32)/255
depth_slice=depth[y0:y0+tile_h,x0:x0+tile_w]
alpha*=np.power(depth_slice,self.cfg.alpha_gamma)
m_slice=mask[y0:y0+tile_h,x0:x0+tile_w]
mask[y0:y0+tile_h,x0:x0+tile_w]=np.maximum(m_slice,alpha*255)
placed+=1
# ── hand dip / face bump ───────────────────────────
m=np.array(mask,dtype=np.float32)
# Hands
for (x0,y0,x1,y1) in self._hands(base):
if random.random()>self.cfg.hand_mask_prob: continue
cx,cy=(x0+x1)/2,(y0+y1)/2
hw,hh=(x1-x0)*self.cfg.hand_expand/2,(y1-y0)*self.cfg.hand_expand/2
ex0,ey0=int(max(0,cx-hw)),int(max(0,cy-hh))
ex1,ey1=int(min(W-1,cx+hw)),int(min(H-1,cy+hh))
h,w=ey1-ey0,ex1-ex0
if h<=0 or w<=0: continue
g=gaussian_bump(h,w)
m[ey0:ey1,ex0:ex1]*=1-g*(1-self.cfg.hand_floor)
# Faces
for (x0,y0,x1,y1) in self._faces(base):
h,w=y1-y0,x1-x0
if h<=0 or w<=0: continue
g=gaussian_bump(h,w)
m[y0:y1,x0:x1]*=1+g*(self.cfg.face_peak-1)
mask_img=Image.fromarray(gaussian_filter(m,self.cfg.blur_sigma).clip(0,255).astype(np.uint8),"L")
out=Path(self.cfg.output_dir); seq=f"image_{self.seq}"
overlay.save(out/f"{seq}.png"); mask_img.save(out/f"{seq}_mask.png")
preview=Image.alpha_composite(dim_base(base,0.3),overlay)
red=Image.merge("RGBA",(mask_img,Image.new("L",(W,H)),Image.new("L",(W,H)),mask_img))
Image.alpha_composite(preview,red).save(out/f"{seq}_preview.png")
Image.fromarray((depth*255).astype(np.uint8),"L").save(out/f"{seq}_depth.png")
with open(out/f"{seq}.txt","w") as f: f.write("\n".join(caps))
mask_img.save(p.parent/f"{p.stem}_mask.png")
if self.cfg.CLEANUP_GENERATED:
for ext in("_depth.png","_preview.png",".txt"): (out/f"{seq}{ext}").unlink(missing_ok=True)
self.seq+=1
def run(self):
imgs=[p for p in Path(self.cfg.images_dir).rglob("*")
if p.suffix.lower() in(".png",".jpg") and not p.stem.endswith(("_mask","_depth"))]
if self.cfg.num_samples: imgs=imgs[:self.cfg.num_samples]
it=tqdm(imgs,desc="Processing",unit="img") if tqdm else imgs
for p in it: self._process(p)
self.depther.unload()
# ───────── MAIN ───────────────────────────────────────────
if __name__=="__main__":
SpiralDepthGen(Config()).run()
print("✔ Done — scaled spirals with hand dips & face bumps generated.")