| """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')]) |
|
|
| |
| 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] |
|
|
| |
| 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]) |
| |
| recent = ckpt_epochs[-5:] |
| print(f'{model_name}: using epochs {recent}') |
|
|
| |
| 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() |
|
|
| |
| 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_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_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() |
|
|
