LAB 11 · 从 RGB 到像素：读出、FCN 与边界

import os, tempfile
from pathlib import Path
# Isolated ephemeral workspace supports detached runs with readonly project storage.
SPATIAL_RUN_DIR = Path(tempfile.mkdtemp(prefix='lecture03-'))
os.chdir(SPATIAL_RUN_DIR)
print('Ephemeral experiment directory:', SPATIAL_RUN_DIR)

import importlib.util, subprocess, sys
from pathlib import Path
# The only download, if needed, is open-source Python packages, not model weights.
if importlib.util.find_spec('torch') is None:
    subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'torch', '--index-url', 'https://download.pytorch.org/whl/cpu'])
missing=[p for p in ['numpy','matplotlib'] if importlib.util.find_spec(p) is None]
if missing: subprocess.check_call([sys.executable,'-m','pip','install',*missing])
import torch, numpy as np
print('Python',sys.version.split()[0], 'torch',torch.__version__, 'numpy',np.__version__, 'device: CPU')

%%writefile lab11_segmentation.py
"""LAB11: sampling warm-up, learned pixel readout, and image-to-mask FCN.
Run: python lab11_segmentation.py --out lab11-results --epochs 20
Dependencies: torch numpy matplotlib. Self-contained; CPU only.
"""
"""Shared transparent geometry and image construction for LAB10/11; no downloads."""
import json, hashlib, platform, sys, time
from pathlib import Path
from datetime import datetime, timezone
import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

def seed_all(seed=7):
    np.random.seed(seed); torch.manual_seed(seed); torch.set_num_threads(2)

def iou_matrix(a,b):
    if len(a)==0 or len(b)==0: return torch.zeros((len(a),len(b)))
    lo=torch.maximum(a[:,None,:2],b[None,:,:2]); hi=torch.minimum(a[:,None,2:],b[None,:,2:])
    inter=(hi-lo).clamp(min=0).prod(-1)
    area_a=(a[:,2:]-a[:,:2]).clamp(min=0).prod(-1)
    area_b=(b[:,2:]-b[:,:2]).clamp(min=0).prod(-1)
    return inter/(area_a[:,None]+area_b[None,:]-inter).clamp(min=1e-8)

def make_scenes(n,seed,size=48,scale=1.0):
    """Independent seed defines original sample, scale changes only object size.
    Boxes xyxy normalized; red rectangle=1, blue ellipse=2; zero to three objects.
    Images are rasterized from latent scene, labels are never a network input.
    """
    rng=np.random.default_rng(seed); xs=[]; ys=[]; masks=[]
    yy,xx=np.mgrid[:size,:size]
    for index in range(n):
        image=rng.uniform(.035,.095,(3,size,size)).astype('float32')
        num=int(rng.choice([0,1,2,3],p=[.18,.30,.36,.16]))
        boxes=[]; labels=[]; mask=np.zeros((size,size),np.int64)
        original=[]
        for obj in range(num):
            for attempt in range(60):
                w,h=rng.integers(9,19,size=2); x=rng.integers(2,size-w-2); y=rng.integers(2,size-h-2)
                box=np.array([x,y,x+w,y+h],np.float32)/size
                if not original or iou_matrix(torch.tensor(np.array(original)),torch.tensor(box[None])).max()<.08: break
            original.append(box); label=int(rng.integers(1,3)); color=rng.uniform(.7,.98)
            cx,cy=(box[0]+box[2])/2,(box[1]+box[3])/2
            sw,sh=(box[2]-box[0])*scale,(box[3]-box[1])*scale
            x1,y1,x2,y2=np.round(np.array([cx-sw/2,cy-sh/2,cx+sw/2,cy+sh/2])*size).astype(int)
            x2=max(x2,x1+2); y2=max(y2,y1+2)
            region=(xx>=x1)&(xx<x2)&(yy>=y1)&(yy<y2)
            if label==2: region &= ((xx+.5-(x1+x2)/2)/((x2-x1)/2))**2+((yy+.5-(y1+y2)/2)/((y2-y1)/2))**2<=1
            image[:,region]*=.2
            image[0 if label==1 else 2,region]=color
            image[1,region]=.15+color*.12
            mask[region]=label
            boxes.append([x1/size,y1/size,x2/size,y2/size]); labels.append(label)
        xs.append(image); masks.append(mask)
        ys.append({'boxes':torch.tensor(boxes,dtype=torch.float32).reshape(-1,4),'labels':torch.tensor(labels,dtype=torch.long),'sample_id':f'{seed}:{index}'})
    return torch.from_numpy(np.stack(xs)),ys,torch.from_numpy(np.stack(masks))

def nms(boxes,scores,threshold=.35):
    order=torch.argsort(scores,descending=True,stable=True); keep=[]
    while len(order):
        i=order[0]; keep.append(i.item()); order=order[1:]
        if len(order): order=order[iou_matrix(boxes[i:i+1],boxes[order])[0]<=threshold]
    return torch.tensor(keep,dtype=torch.long)

def average_precision(predictions,targets,iou_threshold=.5,score_threshold=0):
    """Per-class global score order, max-IoU GT then unused check; all-points AP."""
    result={}; all_tp=all_fp=all_gt=0
    for cls in [1,2]:
        gt={i:t['boxes'][t['labels']==cls] for i,t in enumerate(targets)}
        used={i:set() for i in gt}; total=sum(len(v) for v in gt.values())
        rows=[]
        for i,p in enumerate(predictions):
            for j in range(len(p['scores'])):
                if int(p['labels'][j])==cls and float(p['scores'][j])>=score_threshold:
                    rows.append((float(p['scores'][j]),i,j,p['boxes'][j]))
        rows.sort(key=lambda r:(-r[0],r[1],r[2])); tp=[]; fp=[]
        for score,i,j,b in rows:
            ok=False
            if len(gt[i]):
                overlaps=iou_matrix(b[None],gt[i])[0]; best=int(overlaps.argmax())
                ok=float(overlaps[best])>=iou_threshold and best not in used[i]
                if ok:used[i].add(best)
            tp.append(int(ok));fp.append(int(not ok))
        tp=np.cumsum(tp);fp=np.cumsum(fp)
        recall=tp/max(total,1); precision=tp/np.maximum(tp+fp,1)
        mr=np.r_[0,recall,1];mp=np.r_[0,precision,0]
        mp=np.maximum.accumulate(mp[::-1])[::-1];loc=np.where(mr[1:]!=mr[:-1])[0]
        ap=float(((mr[loc+1]-mr[loc])*mp[loc+1]).sum()) if total else None
        result[str(cls)]={'AP':ap,'GT':total,'TP':int(tp[-1]) if len(tp) else 0,'FP':int(fp[-1]) if len(fp) else 0}
        all_tp+=result[str(cls)]['TP'];all_fp+=result[str(cls)]['FP'];all_gt+=total
    return {'per_class':result,'mAP':float(np.mean([r['AP'] for r in result.values() if r['AP'] is not None])),
            'precision':all_tp/max(all_tp+all_fp,1),'recall':all_tp/max(all_gt,1),'iou_threshold':iou_threshold,'score_threshold':score_threshold}

def provenance(path,started):
    p=Path(path)
    return {'utc_started':started,'utc_completed':datetime.now(timezone.utc).isoformat(),'execution':'CPU, complete script',
            'python':sys.version,'torch':torch.__version__,'numpy':np.__version__,'platform':platform.platform(),'threads':torch.get_num_threads(),
            'code_sha256':hashlib.sha256(p.read_bytes()).hexdigest(),
            'device':'CPU', 'source_file':p.name}

def save_json(path,data):Path(path).write_text(json.dumps(data,indent=2,ensure_ascii=False))

def segmentation_metrics(pred,target):
    valid=target!=255; t=target[valid];p=pred[valid]
    cm=torch.bincount(t*3+p,minlength=9).reshape(3,3).double()
    intersection=cm.diag();union=cm.sum(0)+cm.sum(1)-intersection
    iou=intersection/union.clamp(min=1)
    return {'IoU_background':float(iou[0]),'IoU_red_rectangle':float(iou[1]),'IoU_blue_ellipse':float(iou[2]),
            'mIoU_3_classes':float(iou[union>0].mean()),'pixel_accuracy':float(intersection.sum()/cm.sum()),
            'valid_pixels':int(valid.sum()),'ignored_pixels':int((~valid).sum()),'confusion_matrix_true_rows':cm.int().tolist()}

import argparse

%%writefile -a lab11_segmentation.py
def binary_metrics(pred,target):
    ious=[]
    for cls in [0,1]:
        inter=((pred==cls)&(target==cls)).sum().item();union=((pred==cls)|(target==cls)).sum().item()
        ious.append(inter/union)
    return {'background_IoU':ious[0],'foreground_IoU':ious[1],'mIoU_two_classes':sum(ious)/2,
            'pixel_accuracy':float((pred==target).float().mean()),'error_pixels':int((pred!=target).sum())}

def sampling_warmup(out):
    mask=torch.zeros((8,8),dtype=torch.long);mask[2:6,3:7]=1
    coarse=F.avg_pool2d(mask.float()[None,None],2)
    nearest=F.interpolate(coarse,scale_factor=2,mode='nearest')[0,0]>=.5
    bilinear=F.interpolate(coarse,scale_factor=2,mode='bilinear',align_corners=False)[0,0]>=.5
    nm=binary_metrics(nearest,mask);bm=binary_metrics(bilinear,mask)
    assert abs(nm['foreground_IoU']-2/3)<1e-12
    assert abs(nm['background_IoU']-5/6)<1e-12
    assert abs(nm['mIoU_two_classes']-.75)<1e-12 and nm['error_pixels']==8
    # Bilinear interpolates probabilities/logits; categorical labels use nearest.
    fig,axes=plt.subplots(1,4,figsize=(10,3))
    for ax,data,title in zip(axes,[mask,coarse[0,0],nearest,bilinear],['Known mask','2x2 average','Nearest >=0.5','Bilinear >=0.5']):
        ax.imshow(data,cmap='gray',vmin=0,vmax=1,interpolation='nearest');ax.set_title(title);ax.axis('off')
    fig.tight_layout();fig.savefig(out/'sampling-warmup.png',dpi=150);plt.close(fig)
    return {'nearest':nm,'bilinear_align_corners_false':bm,'identity':'known-mask resampling only; no image model'}

class PixelReadout(nn.Module):
    """A learned softmax regression at every RGB pixel; no spatial receptive field."""
    def __init__(self):super().__init__();self.readout=nn.Conv2d(3,3,1)
    def forward(self,x,mode='bilinear'):return self.readout(x)

class SmallFCN(nn.Module):
    def __init__(self):
        super().__init__()
        self.fine=nn.Sequential(nn.Conv2d(3,12,3,padding=1),nn.ReLU(),nn.Conv2d(12,12,3,padding=1),nn.ReLU())
        self.coarse=nn.Sequential(nn.Conv2d(12,24,3,stride=2,padding=1),nn.ReLU(),nn.Conv2d(24,32,3,stride=2,padding=1),nn.ReLU(),nn.Conv2d(32,32,3,padding=1),nn.ReLU())
        self.head=nn.Conv2d(32,3,1);self.skip=nn.Conv2d(12,3,1)
    def forward(self,x,mode='bilinear'):
        fine=self.fine(x);coarse_logits=self.head(self.coarse(fine))
        kw={'align_corners':False} if mode=='bilinear' else {}
        return F.interpolate(coarse_logits,size=x.shape[-2:],mode=mode,**kw)+self.skip(fine)

def with_ignore(masks):
    masks=masks.clone();masks[:,0,:]=255;masks[:,-1,:]=255;masks[:,:,0]=255;masks[:,:,-1]=255
    # A fixed small region represents missing annotation, not another semantic class.
    masks[:,20:23,20:23]=255
    return masks

%%writefile -a lab11_segmentation.py
def train_model(model,x,target,epochs,lr):
    opt=torch.optim.Adam(model.parameters(),lr=lr);history=[]
    for epoch in range(epochs):
        model.train();losses=[]
        for ids in torch.randperm(len(x)).split(32):
            images=x[ids];labels=target[ids]
            # One geometric decision applies to both image and mask.
            if torch.rand(())<.5:images=images.flip(-1);labels=labels.flip(-1)
            logits=model(images)
            assert logits.shape==(len(ids),3,48,48) and labels.dtype==torch.long
            loss=F.cross_entropy(logits,labels,ignore_index=255)
            opt.zero_grad();loss.backward();opt.step();losses.append(float(loss.detach()))
        history.append(float(np.mean(losses)))
        if epoch==0 or (epoch+1)%5==0:print(type(model).__name__,epoch+1,history[-1],flush=True)
    return history

@torch.no_grad()
def predict(model,x,mode='bilinear'):
    model.eval();return torch.cat([model(b,mode=mode).argmax(1) for b in x.split(32)])

def plot_predictions(images,truth,pred,path,title,indices=None):
    if indices is None:indices=[0,1,2,3]
    palette=np.array([[.04,.04,.05],[.95,.22,.15],[.10,.45,1],[.55,.55,.55]])
    fig,axes=plt.subplots(len(indices),4,figsize=(10,2.6*len(indices)),squeeze=False)
    for row,i in enumerate(indices):
        mask=truth[i].numpy().copy();mask[mask==255]=3
        error=(pred[i]!=truth[i])&(truth[i]!=255)
        data=[images[i].permute(1,2,0).numpy(),palette[mask],palette[pred[i].numpy()],error.numpy()]
        for ax,d,name in zip(axes[row],data,['Input RGB','Target / gray=ignore','Model prediction','Error on valid pixels']):
            ax.imshow(d,cmap='magma',interpolation='nearest');ax.axis('off');ax.set_title(f'{name} #{i}',fontsize=9)
    fig.suptitle(title);fig.tight_layout();fig.savefig(path,dpi=140);plt.close(fig)

%%writefile -a lab11_segmentation.py
def run(out='lab11-results',epochs=20):
    started=datetime.now(timezone.utc).isoformat();start=time.time();out=Path(out);out.mkdir(exist_ok=True,parents=True);seed_all(19)
    warmup=sampling_warmup(out)
    train,_,raw_train=make_scenes(384,7101);test,_,raw_test=make_scenes(128,7901)
    small,_,raw_small=make_scenes(128,7901,scale=.45)
    target=with_ignore(raw_train);test_target=with_ignore(raw_test);small_target=with_ignore(raw_small)
    readout=PixelReadout();pixel_history=train_model(readout,train,target,epochs=25,lr=.03)
    fcn=SmallFCN();random_pred=predict(fcn,test);history=train_model(fcn,train,target,epochs=epochs,lr=.003)
    pred=predict(fcn,test);nearest_pred=predict(fcn,test,mode='nearest');small_pred=predict(fcn,small)
    # Same image/target, blur changes RGB only, preserving geometry and label.
    blurred=F.avg_pool2d(F.pad(test,(2,2,2,2),mode='reflect'),5,stride=1);blur_pred=predict(fcn,blurred)
    # ignore_index check: arbitrarily changing predictions only at ignored pixels leaves metrics identical.
    altered=pred.clone();altered[test_target==255]=(altered[test_target==255]+1)%3
    assert segmentation_metrics(pred,test_target)==segmentation_metrics(altered,test_target)
    # Synchronized flip can be undone on both image and categorical target.
    assert torch.equal(train.flip(-1).flip(-1),train) and torch.equal(target.flip(-1).flip(-1),target)
    results={'provenance':provenance(__file__,started),'config':{'image_shape':[3,48,48],'logits_shape':['B',3,48,48],'target_shape':['B',48,48],
            'target_dtype':'int64','classes':{'0':'background','1':'red rectangle','2':'blue ellipse'},'ignore_label':255,
            'train_count':384,'heldout_count':128,'split_seeds':[7101,7901],'seed':19,'FCN_epochs':epochs,'readout_epochs':25,
            'loss':'cross_entropy on valid pixels, ignore_index=255','optimizer':'Adam FCN lr=.003; pixel readout lr=.03',
            'FCN_parameters':sum(p.numel() for p in fcn.parameters()),'readout_parameters':sum(p.numel() for p in readout.parameters()),
            'upsampling':'12x12 coarse logits ->48x48 bilinear align_corners=False + fine skip logits','augmentation':'paired horizontal flip',
            'data_identity':'generated scenes with color/shape correlated; heldout original scene seeds; no pretrained weights'},
            'sampling_warmup':warmup,'random_FCN':segmentation_metrics(random_pred,test_target),
            'learned_pixel_readout':segmentation_metrics(predict(readout,test),test_target),'trained_FCN':segmentation_metrics(pred,test_target),
            'condition_nearest_logits':segmentation_metrics(nearest_pred,test_target),'condition_scale_045':segmentation_metrics(small_pred,small_target),
            'condition_RGB_blur5':segmentation_metrics(blur_pred,test_target),'readout_RGB_blur5':segmentation_metrics(predict(readout,blurred),test_target),'readout_scale_045':segmentation_metrics(predict(readout,small),small_target),'history_FCN':history,'history_pixel_readout':pixel_history,
            'ignore_metric_invariance_test':'passed','paired_flip_involution_test':'passed','sampling_known_values_tests':'passed','seconds':time.time()-start}
    counts=((pred!=test_target)&(test_target!=255)).flatten(1).sum(1);worst=counts.argsort(descending=True)[:4].tolist()
    results['failure_examples']=[{'sample_index':i,'wrong_valid_pixels':int(counts[i])} for i in worst]
    torch.save({'state_dict':fcn.state_dict(),'config':results['config']},out/'fcn.pt');torch.save(readout.state_dict(),out/'pixel-readout.pt')
    np.savez_compressed(out/'heldout-input-target-predictions.npz',images=test.numpy(),target=test_target.numpy(),prediction=pred.numpy(),
                       half_scale_images=small.numpy(),half_scale_target=small_target.numpy(),half_scale_prediction=small_pred.numpy(),blur_images=blurred.numpy(),blur_prediction=blur_pred.numpy())
    plot_predictions(test,test_target,pred,out/'heldout-segmentation.png','FCN heldout image-to-mask predictions')
    plot_predictions(test,test_target,pred,out/'failure-segmentation.png','Four largest valid-pixel error counts',worst)
    plot_predictions(small,small_target,small_pred,out/'small-object-segmentation.png','Same latent scenes: object dimensions x0.45')
    plot_predictions(blurred,test_target,blur_pred,out/'blur-segmentation.png','Same geometry, RGB blurred with 5x5 average')
    fig,ax=plt.subplots(figsize=(7,4));ax.plot(history,label='FCN');ax.plot(pixel_history,label='pixel softmax readout');ax.set(xlabel='epoch (zero based)',ylabel='valid-pixel CE');ax.legend();fig.tight_layout();fig.savefig(out/'training-curve.png',dpi=150);plt.close(fig)
    save_json(out/'run.json',results);print(json.dumps({k:results[k]['mIoU_3_classes'] for k in ['random_FCN','learned_pixel_readout','trained_FCN','condition_scale_045','condition_RGB_blur5']},indent=2),flush=True)
    return results

if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('--out',default='lab11-results');p.add_argument('--epochs',type=int,default=20);a=p.parse_args();run(a.out,a.epochs)

import subprocess,sys,json
subprocess.run([sys.executable,'lab11_segmentation.py','--out','lab11-results-rerun','--epochs','20'],check=True)
R=json.loads(Path('lab11-results-rerun/run.json').read_text())
print('sampling:',R['sampling_warmup'])
for key in ['random_FCN','learned_pixel_readout','trained_FCN','condition_nearest_logits','condition_scale_045','condition_RGB_blur5','readout_RGB_blur5']:
    r=R[key];print(key,{k:round(v,5) for k,v in r.items() if 'IoU' in k or k=='pixel_accuracy'})
print(R['provenance'])
from IPython.display import Image,display
for filename in ['sampling-warmup.png','heldout-segmentation.png','failure-segmentation.png','small-object-segmentation.png','blur-segmentation.png']:
    display(Image(filename='lab11-results-rerun/'+filename))