| """Procedural object-labelled skins. Labels describe authored objects, not RGB rules.""" |
| import numpy as np |
| import torch |
| from torch.utils.data import Dataset |
| from SkingToolkit.dense_uv_parser.skin_dataset import load_skin |
|
|
|
|
| def make_accessory_skin(original, seed, return_components=False): |
| rng = np.random.default_rng(seed) |
| uv = original.permute(1, 2, 0).numpy().copy() |
| objects = np.zeros((64, 64), dtype=np.int64) |
| components = np.zeros((64,64),dtype=np.int64) |
| |
| uv[:16, :32] = 0 |
| uv[:16, 32:] = 0 |
| faces = [(8, 8), (24, 8), (0, 8), (16, 8), (8, 0), (16, 0)] |
| skin = np.array(rng.choice([[.87,.63,.43],[.45,.28,.19],[.65,.44,.30], |
| [.95,.78,.62],[.32,.20,.13],[.74,.62,.49]])) |
| if rng.random() < .12: |
| skin = rng.uniform(.2, .9, 3) |
| hair = rng.uniform(.03, .30, 3) if rng.random() < .7 else rng.uniform(.1, .85, 3) |
| for fi, (x, y) in enumerate(faces): |
| tex = np.clip(skin[None,None] + rng.normal(0,.025,(8,8,1)), 0, 1) |
| if fi in (1,4): tex[:] = hair + rng.normal(0,.025,(8,8,1)) |
| if fi in (0,2,3): tex[:rng.integers(1,3)] = hair |
| uv[y:y+8, x:x+8, :3] = np.clip(tex,0,1) |
| uv[y:y+8, x:x+8, 3] = 1 |
| |
| for x in (9, 13): |
| uv[11:12, x:x+2, :3] = rng.uniform(.01,.18,3) |
| uv[14,11:13,:3] = rng.uniform(.05,.5,3) |
| if rng.random() < .4: |
| uv[10,9:11,:3] = hair |
| uv[10,13:15,:3] = hair |
|
|
| def put(face, mask, color, identity, layer=1, component=0): |
| x, y = faces[face]; x += 32*layer |
| patch = uv[y:y+8, x:x+8] |
| patch[mask, :3] = color[mask] if np.asarray(color).ndim == 3 else color |
| patch[mask, 3] = 1 |
| objects[y:y+8, x:x+8][mask] = identity |
| components[y:y+8, x:x+8][mask] = component |
|
|
| yy, xx = np.mgrid[:8,:8] |
| negative = rng.random() < .18 |
| if not negative and rng.random() < .78: |
| top = int(rng.integers(2,4)); height = int(rng.integers(3,5)) |
| frame_color = rng.uniform(.025,.23,3) if rng.random()<.70 else rng.uniform(.01,.85,3) |
| lens_color = rng.uniform(.05,.95,3) |
| frame = np.zeros((8,8),bool); lenses = np.zeros_like(frame) |
| |
| visor = rng.random()<.20 |
| boxes = [(0,3),(4,7)] if visor else [(0,2),(5,7)] |
| hollow = rng.random()<.30 |
| rimless_top = rng.random()<.45 and not hollow |
| for box_index,(left, right) in enumerate(boxes): |
| rectangle = (xx>=left)&(xx<=right)&(yy>=top)&(yy<top+height) |
| edge = rectangle & ((xx==left)|(xx==right)|(yy==top)|(yy==top+height-1)) |
| if rimless_top: |
| edge &= (yy!=top) |
| |
| if not visor and not hollow and rng.random()<.65: |
| inner_x=right if box_index==0 else left |
| edge &= ~((xx==inner_x)&(yy<top+height-1)) |
| if rng.random() < .2: |
| edge &= ~(((xx==left)|(xx==right))&(yy==top+height-1)) |
| frame |= edge |
| lenses |= rectangle & ~edge |
| bridge_y = min(7,top+int(rng.integers(0,2))) |
| frame[bridge_y,2:6] = True |
| lenses[frame] = False |
| if hollow: lenses[:] = False |
| color = np.broadcast_to(frame_color,(8,8,3)).copy() |
| color[lenses] = lens_color |
| |
| color[lenses] *= rng.uniform(.65,1.05,(8,8,1))[lenses] |
| put(0, frame|lenses, np.clip(color,0,1), 1) |
| for face in (2,3): |
| length = int(rng.integers(3,9)) |
| temple = (yy==bridge_y)&(xx<length) |
| put(face,temple,frame_color,1) |
| kind = int(rng.choice([0,2,3], p=[.35,.40,.25])) if not negative else 0 |
| if kind==2: |
| |
| |
| color = rng.uniform(.02,.17,3) if rng.random()<.65 else rng.uniform(.05,.85,3) |
| depth = int(rng.integers(2,6)) |
| stepped = rng.random()<.70 |
| variant = int(rng.choice(3,p=[.75,.15,.10])) |
| band_color = rng.uniform(.04,.95,3) |
| if rng.random()<.3: band_color=np.array([rng.uniform(.5,.95),rng.uniform(.02,.12),rng.uniform(.02,.12)]) |
| if rng.random()<.2: band_color=color*rng.uniform(.7,1.2) |
| has_band = rng.random()<.85 |
| crown_layer=0 if stepped else 1 |
| crown_component=1 if stepped else 3 |
| band_component=2 if stepped else 4 |
| put(4,np.ones((8,8),bool),color,0 if stepped else 2,layer=crown_layer,component=crown_component) |
| for face in (0,1,2,3): |
| tex=np.clip(np.broadcast_to(color,(8,8,3))+rng.normal(0,.02,(8,8,1)),0,1) |
| crown=yy<depth-1 |
| put(face,crown,tex,0 if stepped else 2,layer=crown_layer,component=crown_component) |
| if has_band: |
| put(face,yy==depth-2,np.clip(band_color,0,1),0 if stepped else 2,layer=crown_layer,component=band_component) |
| brim=yy==depth-1 |
| if variant==1 and face!=0: brim[:]=False |
| if variant==2 and face in (1,3): brim &= xx<5 |
| put(face,brim,color,2,component=5) |
| elif kind==3: |
| put(4,np.ones((8,8),bool),hair,3) |
| for face in (0,1,2,3): |
| depth=int(rng.integers(1,4)) if face==0 else int(rng.integers(2,7)) |
| mask=(yy<depth)|((yy==depth)&(rng.random((8,8))>.45)) |
| tex=np.clip(np.broadcast_to(hair,(8,8,3))+rng.normal(0,.025,(8,8,1)),0,1) |
| put(face,mask,tex,3) |
| if return_components: |
| return torch.from_numpy(uv).permute(2,0,1).float(),torch.from_numpy(objects),torch.from_numpy(components) |
| return torch.from_numpy(uv).permute(2,0,1).float(), torch.from_numpy(objects) |
|
|
|
|
| class AccessoryDataset(Dataset): |
| def __init__(self, paths, count, seed): |
| self.paths, self.count, self.seed, self.epoch = list(paths), count, seed, 0 |
|
|
| def __len__(self): return self.count |
|
|
| def __getitem__(self, index): |
| seed = self.seed + index*104729 + self.epoch*self.count*104729 |
| source = load_skin(self.paths[(index + self.epoch*self.count) % len(self.paths)]) |
| uv, objects, components = make_accessory_skin(source, seed, return_components=True) |
| return {'uv':uv,'objects':objects,'components':components,'seed':seed} |
|
|
|
|
| def render_accessories(uv, objects, renderer, views, components=None): |
| from SkingToolkit.dense_uv_parser.utils import build_dense_parser_batch |
| images, labels, foregrounds, parts = [], [], [], [] |
| for view in views: |
| rendered, target = build_dense_parser_batch(uv, renderer, view) |
| xy = (target['uv']*63).round().long().clamp(0,63) |
| flat = xy[:,1]*64+xy[:,0] |
| identity = objects.flatten(1).gather(1,flat.flatten(1)).reshape_as(flat) |
| identity = identity.masked_fill((target['layer'] != 1) | (target['foreground'][:,0] < .5), 0) |
| images.append(rendered);labels.append(identity);foregrounds.append(target['foreground'][:,0].bool()) |
| if components is not None: |
| part=components.flatten(1).gather(1,flat.flatten(1)).reshape_as(flat) |
| parts.append(part.masked_fill(target['foreground'][:,0]<.5,0)) |
| b = uv.shape[0] |
| result=(torch.stack(images,1).flatten(0,1),torch.stack(labels,1).flatten(0,1),torch.stack(foregrounds,1).flatten(0,1)) |
| return result+(torch.stack(parts,1).flatten(0,1),) if components is not None else result |
|
|
