读源码时跟踪三种集合 · 实现代码

%%writefile lab10_detector.py
"""LAB10: train and evaluate an image-to-box anchor detector, CPU, no downloads.
Run: python lab10_detector.py --out lab10-results --epochs 28
Dependencies: torch numpy matplotlib. This script is self-contained.
"""
"""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