| |
| """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 |
|
|
| 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 |
|
|
| |
| 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") |
|
|
| |
| 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 |
| |
| 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 |
|
|
| 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 |
| |
| if len(wbf_b) == 0: continue |
|
|
| |
| 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() |
|
|
