"""Checkpoint Ensemble: fuse predictions from different training stages. ==================================================================== Same model at epoch 80/90/100/110/120 has different error patterns. WBF across 5 checkpoints = free diversity, zero training cost. Also supports per-epoch SWA (Stochastic Weight Averaging) checkpoint generation. """ import sys, os, json, gc import numpy as np from PIL import Image, ImageEnhance from tqdm import tqdm from collections import defaultdict PROJECT_DIR = '/home/user/goat' os.chdir(PROJECT_DIR) sys.path.insert(0, PROJECT_DIR) import torch from ultralytics import YOLO def compute_iou(b1, b2): x1,y1 = max(b1[0],b2[0]), max(b1[1],b2[1]) x2,y2 = min(b1[2],b2[2]), min(b1[3],b2[3]) inter = max(0,x2-x1)*max(0,y2-y1) a1 = (b1[2]-b1[0])*(b1[3]-b1[1]); a2 = (b2[2]-b2[0])*(b2[3]-b2[1]) return inter/(a1+a2-inter+1e-8) def wbf(boxes_list, scores_list, iou_thr=0.55): if not boxes_list or all(len(b)==0 for b in boxes_list): return np.array([]), np.array([]) all_boxes, all_scores = [], [] for boxes, scores in zip(boxes_list, scores_list): all_boxes.extend(boxes); all_scores.extend(scores) if not all_boxes: return np.array([]), np.array([]) all_boxes = np.array(all_boxes); all_scores = np.array(all_scores) order = np.argsort(-all_scores) all_boxes, all_scores = all_boxes[order], all_scores[order] clusters, used = [], np.zeros(len(all_boxes), dtype=bool) for i in range(len(all_boxes)): if used[i]: continue cluster = [(all_boxes[i], all_scores[i])]; used[i] = True for j in range(i+1, len(all_boxes)): if used[j]: continue tw = sum(s for _, s in cluster) center = sum(b*s/tw for b, s in cluster) if compute_iou(center.tolist(), all_boxes[j].tolist()) > iou_thr: cluster.append((all_boxes[j], all_scores[j])); used[j] = True clusters.append(cluster) result_boxes, result_scores = [], [] for cl in clusters: tw = sum(s for _, s in cl) avg_b = sum(b*s/tw for b, s in cl); result_boxes.append(avg_b); result_scores.append(tw) return np.array(result_boxes), np.array(result_scores) def main(): val_img_dir = 'Data/Detection_dataset/images/val' val_label_dir = 'Data/Detection_dataset/labels/val' val_files = sorted([f for f in os.listdir(val_img_dir) if f.endswith('.jpg')]) # Test with v15_seed_333 (has epoch checkpoints saved) exp_name = 'v15_seed_333' exp_dir = f'runs/detect/Detection_experiments/{exp_name}/weights' ckpts = sorted([f for f in os.listdir(exp_dir) if f.startswith('epoch') and f.endswith('.pt')]) ckpt_epochs = [int(f.replace('epoch','').replace('.pt','')) for f in ckpts] # Also try v6_1 checkpoints ckpt_dirs = [ ('v6_1_s_refined', 'runs/detect/Detection_experiments/v6_1_s_refined/weights'), ('v15_seed_333', 'runs/detect/Detection_experiments/v15_seed_333/weights'), ] all_eval_data = [] for model_name, ckpt_dir in ckpt_dirs: ckpts = sorted([f for f in os.listdir(ckpt_dir) if f.startswith('epoch') and f.endswith('.pt')]) if len(ckpts) < 3: continue ckpt_epochs = sorted([int(f.replace('epoch','').replace('.pt','')) for f in ckpts]) # Take last 5 checkpoints recent = ckpt_epochs[-5:] print(f'{model_name}: using epochs {recent}') # Run predictions for each checkpoint ckpt_preds = {} for ep in tqdm(recent, desc=model_name): path = os.path.join(ckpt_dir, f'epoch{ep}.pt') if not os.path.exists(path): continue m = YOLO(path) img_preds = [] for img_file in tqdm(val_files, desc=f' ep{ep}', leave=False): img = Image.open(os.path.join(val_img_dir, img_file)) r = m.predict(img, imgsz=1536, conf=0.25, iou=0.7, max_det=100, verbose=False) if r and len(r[0].boxes): img_preds.append((r[0].boxes.xyxy.cpu().numpy(), r[0].boxes.conf.cpu().numpy())) else: img_preds.append((np.array([]), np.array([]))) ckpt_preds[ep] = img_preds del m; gc.collect(); torch.cuda.empty_cache() # Evaluate: single best checkpoint, and WBF of all 5 iou_thrs = [round(0.5+i*0.05,2) for i in range(10)] def eval_preds(name, preds_dict, val_files): tp = {t:0 for t in iou_thrs}; total_gt = 0 for idx, img_file in enumerate(val_files): img = Image.open(os.path.join(val_img_dir, img_file)) gt_boxes = [] lf = img_file.replace('.jpg','.txt') with open(os.path.join(val_label_dir, lf)) as f: for line in f: p = line.strip().split() if len(p)>=5: cx,cy,w,h = [float(x) for x in p[1:5]] gt_boxes.append([(cx-w/2)*img.size[0],(cy-h/2)*img.size[1],(cx+w/2)*img.size[0],(cy+h/2)*img.size[1]]) total_gt += len(gt_boxes) if not gt_boxes: continue preds = preds_dict[idx] for t in iou_thrs: matched = set() for pb in preds: if len(pb)==0: continue best_iou, best_gi = 0, -1 for gi, gb in enumerate(gt_boxes): if gi in matched: continue iou = compute_iou(pb.tolist(), gb) if iou > best_iou: best_iou = iou; best_gi = gi if best_iou >= t and best_gi >= 0: tp[t] += 1; matched.add(best_gi) recalls = [tp[t]/total_gt for t in iou_thrs] mAP = np.mean(recalls) return mAP, recalls[5] # Best single checkpoint best_ep, best_mAP = 0, 0 for ep, preds in ckpt_preds.items(): mAP, r75 = eval_preds(f'ep{ep}', [p[0] for p in preds], val_files) if mAP > best_mAP: best_mAP = mAP; best_ep = ep print(f' Best single cp: ep{best_ep} mAP={best_mAP:.4f}') # WBF of all 5 wbf_preds = [] for idx in range(len(val_files)): bl, sl = [], [] for ep, preds in ckpt_preds.items(): b, s = preds[idx] if len(b) > 0: bl.append(b); sl.append(s) wbf_preds.append(wbf(bl, sl)[0]) mAP_wbf, r75_wbf = eval_preds('WBF(5cp)', wbf_preds, val_files) print(f' WBF of 5 cps: mAP={mAP_wbf:.4f} (+{mAP_wbf-best_mAP:+.4f} vs best single)') all_eval_data.append({'model': model_name, 'best_cp': best_mAP, 'wbf_5cp': mAP_wbf}) with open('logs/checkpoint_ensemble.json', 'w') as f: json.dump(all_eval_data, f, indent=2) print('\nSaved to logs/checkpoint_ensemble.json') if __name__ == '__main__': main()