File size: 7,020 Bytes
6a5bb7e | 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 152 153 154 155 156 157 158 159 160 | """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()
|