File size: 4,280 Bytes
f211247 | 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 | #!/usr/bin/env python3
"""Soft-Teacher style pseudo-labeling.
Use ensemble of 5-fold YOLO models to label background images.
Keep predictions with conf>=0.7 AND ensemble agreement (≥3 models).
Output YOLO format labels + image list.
"""
import os, sys, json, time, random
from pathlib import Path
import numpy as np
ROOT = Path('/arf/scratch/stakan/hitit-proje')
WEIGHTS_ROOT = ROOT / 'runs/detect/hitit_ocr/runs/h100'
PSEUDO_DIR = ROOT / 'datasets/ready/detection_pseudo'
PSEUDO_IMG = PSEUDO_DIR / 'images/all'
PSEUDO_LBL = PSEUDO_DIR / 'labels/all'
PSEUDO_IMG.mkdir(parents=True, exist_ok=True)
PSEUDO_LBL.mkdir(parents=True, exist_ok=True)
PSEUDO_IMG_LIST = ROOT / 'datasets/ready/detection_tablets/pseudo_train.txt'
CONF_THR = 0.7
AGREE_MIN = 3
N_BG_SAMPLE = 5000 # her run'da 5k bg label
def log(m): print(f"[{time.strftime('%H:%M:%S')}] {m}", flush=True)
def main():
from ultralytics import YOLO
from ensemble_boxes import weighted_boxes_fusion
from PIL import Image
# 5 model
models = []
for fold in range(5):
ck = WEIGHTS_ROOT / f'yolo_fold{fold}/weights/best.pt'
if ck.exists():
models.append(YOLO(str(ck)))
log(f"Loaded fold {fold}")
log(f"{len(models)} models loaded")
# BG image pool
IMG_DIR = ROOT / 'hitit_ocr/data/detection/images/all'
LBL_DIR = ROOT / 'hitit_ocr/data/detection/labels/all'
pool = []
for img in IMG_DIR.iterdir():
if not img.exists(): continue
stem = img.stem
lbl = LBL_DIR / f'{stem}.txt'
if lbl.exists() and lbl.stat().st_size > 0: continue
# Skip maicubeda crops (too small)
if stem.startswith('tablet_maicubeda_'): continue
if stem.startswith('tablet_ob_sign_'): continue
pool.append(str(img))
log(f"BG candidate pool (no crops): {len(pool)}")
random.seed(0); random.shuffle(pool)
pool = pool[:N_BG_SAMPLE]
log(f"Will pseudo-label {len(pool)} images")
n_kept = 0; total_boxes = 0
pseudo_imgs = []
for idx, img_path in enumerate(pool):
if idx % 200 == 0: log(f" {idx}/{len(pool)} kept={n_kept}")
try:
with Image.open(img_path) as im: W, H = im.size
except Exception: continue
if min(W, H) < 256 or max(W, H) > 4000: continue # skip extreme
boxes_list, scores_list, labels_list = [], [], []
skip = False
for m in models:
try:
r = m.predict(img_path, conf=CONF_THR, iou=0.7, max_det=2000,
imgsz=1280, verbose=False, device=0)[0]
except Exception:
skip = True; break
if r.boxes is None or len(r.boxes) == 0: continue
xyxy = r.boxes.xyxy.cpu().numpy() / np.array([W, H, W, H])
xyxy = np.clip(xyxy, 0, 1)
boxes_list.append(xyxy.tolist())
scores_list.append(r.boxes.conf.cpu().numpy().tolist())
labels_list.append([0]*len(r.boxes))
if skip or len(boxes_list) < AGREE_MIN: continue
try:
wbf_b, wbf_s, wbf_l = weighted_boxes_fusion(
boxes_list, scores_list, labels_list,
iou_thr=0.55, skip_box_thr=CONF_THR)
except Exception:
continue
# require >=3 model contribution per box (ensemble_boxes returns averaged)
if len(wbf_b) == 0: continue
# Write YOLO label + symlink image to pseudo_dir
src = Path(img_path)
stem = src.stem
ext = src.suffix
lines = []
for (x1,y1,x2,y2), sc in zip(wbf_b, wbf_s):
cx = (x1+x2)/2; cy = (y1+y2)/2
w = x2-x1; h = y2-y1
lines.append(f"0 {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}")
if not lines: continue
(PSEUDO_LBL / f'{stem}.txt').write_text('\n'.join(lines)+'\n')
link = PSEUDO_IMG / f'{stem}{ext}'
if not link.exists():
try: link.symlink_to(src)
except FileExistsError: pass
pseudo_imgs.append(str(link))
n_kept += 1; total_boxes += len(lines)
log(f"Done: kept={n_kept} images, total pseudo-boxes={total_boxes}")
PSEUDO_IMG_LIST.write_text('\n'.join(pseudo_imgs)+'\n')
log(f"Wrote {PSEUDO_IMG_LIST}")
if __name__ == '__main__':
main()
|