Script+Experiment; auto-masking with depth and spiral masks.

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/211f31f8d901dadd1025398fa1a603b8

To 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.")