File size: 7,831 Bytes
4ae53b1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | """Round 4: matched multiscale structure and DDColor distillation experiments."""
import os,sys,json,time,random,hashlib,traceback,shutil
from pathlib import Path
import numpy as np
import cv2
from PIL import Image,ImageDraw
from skimage import data as sample_data
from skimage.color import rgb2lab,lab2rgb
import pyarrow as pa
import pyarrow.parquet as pq
import torch
import torch.nn.functional as F
from torch.utils.data import Dataset,DataLoader
from huggingface_hub import HfApi,hf_hub_download,PyTorchModelHubMixin
from model import load_model,save_model
from data import decode_image,ColorBins
from metrics import soft_ce,per_image_metrics
from spatial import guided_chroma
OUT=Path('round6');OUT.mkdir(exist_ok=True)
DURABLE=None
START=time.monotonic();DEADLINE=START+3.5*3600
DEVICE='cuda';SEED=409;STEPS=3000
BASE_REV='704fa80d792c3d759db91daa00b2dcfe6f0f6412'
BASE_HASH='0e4c417375684a044860f8af3ac3a2fb44e1a5729254ca33abda758ea71aea6e'
TEACHER_ID='piddnad/ddcolor_artistic'
torch.set_num_threads(6)
def deadline():
if time.monotonic()>DEADLINE:raise TimeoutError('Graceful export before remote hard timeout')
def gray_arrays(im,style='gray'):
im=im.resize((256,256),Image.Resampling.BILINEAR)
rgb=np.asarray(im,np.float32)/255
lab=rgb2lab(rgb).astype(np.float32)
if style=='film':
g=(rgb*np.array([.42,.45,.13],np.float32)).sum(-1)
g=np.clip((g**1.15-.5)*.8+.5,0,1)
else:g=np.asarray(im.convert('L'),np.float32)/255
# Feed the same neutral sRGB input to teacher and student.
gray=np.repeat(g[...,None],3,-1).astype(np.float32)
L=rgb2lab(gray)[...,0].astype(np.float32)
return (L[None]/50-1).copy(),lab[...,1:].transpose(2,0,1).copy(),gray
class Photos(Dataset):
def __init__(self,table,ids,bins=None,cache=None,augment=False,style='gray'):
self.table,self.ids,self.bins,self.cache,self.augment,self.style=table,ids,bins,cache,augment,style
def __len__(self):return len(self.ids)
def __getitem__(self,k):
L,ab,gray=gray_arrays(decode_image(self.table,self.ids[k]),self.style)
teach=np.array(self.cache[k],np.float32) if self.cache is not None else np.zeros((2,64,64),np.float32)
if self.augment:
if random.random()<.5:L=L[...,::-1].copy();ab=ab[...,::-1].copy();teach=teach[...,::-1].copy()
if random.random()<.3:L=np.clip(L*random.uniform(.85,1.15)+random.uniform(-.08,.08),-1,1)
if self.bins is not None:
idx,w=self.bins.encode(ab.transpose(1,2,0));return torch.from_numpy(L),torch.from_numpy(ab),torch.from_numpy(idx),torch.from_numpy(w),torch.from_numpy(teach)
return torch.from_numpy(L),torch.from_numpy(ab),torch.from_numpy(gray.transpose(2,0,1)),self.ids[k]
def patch_excess(p,t):
"""Excess color differences across 4–64px regions, in target-flat areas."""
values=[]
for scale in [4,16]:
a=F.avg_pool2d(p,scale);b=F.avg_pool2d(t,scale)
for offset in [1,4]:
for dim in [-1,-2]:
x=a.narrow(dim,offset,a.shape[dim]-offset)-a.narrow(dim,0,a.shape[dim]-offset)
y=b.narrow(dim,offset,b.shape[dim]-offset)-b.narrow(dim,0,b.shape[dim]-offset)
dx=x.norm(dim=1);dy=y.norm(dim=1);mask=(dy<3).float()
values.append(((dx-dy-1).relu()*mask).sum((1,2))/mask.sum((1,2)).clamp_min(1))
return sum(values)/len(values)
def structural_loss(p,t):
terms=[]
for scale in [4,16]:
a=F.avg_pool2d(p,scale);b=F.avg_pool2d(t,scale)
for offset in [1,4]:
for dim in [-1,-2]:
x=a.narrow(dim,offset,a.shape[dim]-offset)-a.narrow(dim,0,a.shape[dim]-offset)
y=b.narrow(dim,offset,b.shape[dim]-offset)-b.narrow(dim,0,b.shape[dim]-offset)
terms.append(F.smooth_l1_loss(x/10,y/10,beta=.3))
return sum(terms)/len(terms)
def additional_losses(p,t,teacher):
small=F.avg_pool2d(p,4);target=F.avg_pool2d(t,4)
# Teacher is fallible; reduce guidance when it conflicts strongly with ground truth.
confidence=torch.exp(-(teacher-target).norm(dim=1,keepdim=True)/30).detach()
kd=(F.smooth_l1_loss(small/20,teacher/20,beta=.5,reduction='none')*confidence).mean()
neutral=(t.norm(dim=1)<3).float()
neutral_loss=(p.norm(dim=1)*neutral).sum()/neutral.sum().clamp_min(1)/20
return structural_loss(p,t),kd,neutral_loss
def measure(L,p,t,raw=None):
r=per_image_metrics(L,p,t);C=p.norm(dim=1);T=t.norm(dim=1)
for name,mask,event in [('missed_color',T>12,C<5),('neutral_spill',T<3,C>10)]:
r[name]=(mask&event).sum((1,2))/mask.sum((1,2)).clamp_min(1)
r['color_coverage']=(C>10).float().mean((1,2));r['patch_excess']=patch_excess(p,t)
r['raw_patch_excess']=patch_excess(p if raw is None else raw,t)
return r
@torch.inference_mode()
def teacher_predict(teacher,gray):
gray=F.interpolate(gray.to(DEVICE),size=(512,512),mode='bilinear',align_corners=False)
# Full precision is deliberate: spectral normalization / attention are not assumed BF16-safe.
out=teacher(gray).float()
assert out.shape[1]==2 and torch.isfinite(out).all()
return F.interpolate(out,size=(256,256),mode='bilinear',align_corners=False)
@torch.inference_mode()
def evaluate(model,table,ids,is_teacher=False):
result={};model.eval()
for style in ['gray','film']:
rows=[]
for L,t,gray,indices in DataLoader(Photos(table,ids,style=style),batch_size=2 if is_teacher else 8,num_workers=4):
deadline();L,t=L.to(DEVICE),t.to(DEVICE)
raw=teacher_predict(model,gray) if is_teacher else model.decode(model(L),.38)
pred=raw if is_teacher else guided_chroma(L,raw,8)
metrics=measure(L,pred,t,raw)
for j,index in enumerate(indices):rows.append({'index':int(index)}|{k:float(v[j]) for k,v in metrics.items()})
keep=[x for x in rows if x['target_chroma']>=5]
result[style]={'n_total':len(rows),'n_color':len(keep),'summary':{k:float(np.mean([x[k] for x in keep])) for k in keep[0] if k!='index'},'per_image':rows}
return result
def selection(result,baseline):
scores=[];eligible=True
for style in ['gray','film']:
s=result[style]['summary'];b=baseline[style]['summary']
eligible &= (s['ab_error']<=b['ab_error']*1.05 and s['neutral_spill']<=b['neutral_spill']+.01
and s['missed_color']<=b['missed_color']+.01 and s['color_coverage']>=.95*b['color_coverage']
and s['patch_excess']<.85*b['patch_excess'] and s['raw_patch_excess']<.9*b['raw_patch_excess'])
scores.append(s['patch_excess']/max(b['patch_excess'],1e-6)+.3*s['neutral_spill']+.2*s['ab_error']/b['ab_error'])
return float(np.mean(scores)),bool(eligible)
@torch.inference_mode()
def probes(model,tag,teacher=False):
names=['astronaut','coffee','chelsea','rocket','camera','coins','moon'];images=[]
for name in names:
original=Image.fromarray(getattr(sample_data,name)()).convert('RGB')
L,_,gray=gray_arrays(original);x=torch.from_numpy(L)[None].to(DEVICE)
ab=teacher_predict(model,torch.from_numpy(gray.transpose(2,0,1))[None]) if teacher else guided_chroma(x,model.decode(model(x),.38),8)
lab=np.concatenate([(L[0]*50+50)[...,None],ab[0].cpu().numpy().transpose(1,2,0)],-1)
rgb=np.uint8(np.clip(lab2rgb(lab),0,1)*255)
images.append((name,Image.fromarray(np.uint8(gray*255)),Image.fromarray(rgb)))
canvas=Image.new('RGB',(512,len(names)*280),'white');d=ImageDraw.Draw(canvas)
for i,(name,gray,col) in enumerate(images):
canvas.paste(gray,(0,i*280+24));canvas.paste(col,(256,i*280+24));d.text((4,i*280+4),name+' input | '+tag,fill='black')
canvas.save(OUT/(tag+'.jpg'),quality=90)
def write(name,obj):
(OUT/name).write_text(json.dumps(obj,indent=2));return obj
|