code/src/preprocessing/uncertainty_dinov3.py

#!/usr/bin/env python3
"""DINOv3 tabanlı uncertainty score.

Yaklaşım: DINOv3 backbone'u ile feature çıkar, her unified_label'ın centroid'ini hesapla,
her kayıt için (1 - max_similarity_to_own_centroid) → uncertainty.
Hard cases (düşük similarity) = yüksek uncertainty = tekrar annotate önerilir.
"""
import json, os, argparse, time
from pathlib import Path
import numpy as np

ROOT = Path("/arf/scratch/stakan/hitit-proje")
SOURCES = ROOT / "datasets" / "sources"

def main():
    import torch
    from PIL import Image
    from collections import defaultdict
    
    ap = argparse.ArgumentParser()
    ap.add_argument('--batch', type=int, default=64)
    ap.add_argument('--model', default='dinov2_vitb14')  # DINOv3 eğer torch hub'da varsa dinov3_vitb16
    ap.add_argument('--sample', type=int, default=0)
    args = ap.parse_args()
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Device: {device}", flush=True)
    
    # DINOv3 (Meta) — Py3.11 + ai-tools-kolyoz-1.0 env'i ile çalışır.
    # Fallbacks: DINOv2 → timm DINOv2 (en kötü senaryo).
    try:
        model = torch.hub.load('facebookresearch/dinov3', 'dinov3_vitb16',
                                trust_repo=True).eval().to(device)
        emb_dim = 768
        input_size = 224
        model_name = "dinov3_vitb16"
    except Exception as e:
        print(f"DINOv3 yüklenemedi ({type(e).__name__}: {str(e)[:80]}), DINOv2'ye düşüldü", flush=True)
        try:
            model = torch.hub.load('facebookresearch/dinov2', 'dinov2_vitb14',
                                    trust_repo=True).eval().to(device)
            emb_dim = 768
            input_size = 224
            model_name = "dinov2_vitb14"
        except Exception as e2:
            print(f"DINOv2 de yüklenemedi ({e2}), timm DINOv2'ye düşüldü", flush=True)
            import timm
            model = timm.create_model('vit_base_patch14_dinov2.lvd142m',
                                       pretrained=True, num_classes=0).eval().to(device)
            emb_dim = 768
            input_size = 518
            model_name = "dinov2_vitb14_timm"
    
    from torchvision import transforms
    tf = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(input_size),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),
    ])
    
    # Topla: sadece classification
    items = []
    for d in sorted(SOURCES.iterdir()):
        if not d.is_dir(): continue
        mp = d / "manifest_classification.jsonl"
        if not mp.exists(): continue
        with open(mp) as f:
            for line in f:
                r = json.loads(line)
                p = r.get('path')
                ul = r.get('unified_label')
                if p and ul and r.get('integrity_ok') is True and os.path.exists(p):
                    items.append((r['id'], p, ul))
    if args.sample:
        import random; random.seed(42)
        items = random.sample(items, min(args.sample, len(items)))
    print(f"Feature çıkarılacak: {len(items):,}", flush=True)
    
    # Feature extraction
    features = np.zeros((len(items), emb_dim), dtype=np.float32)
    labels_list = [None] * len(items)
    ids_list = [None] * len(items)
    
    t0 = time.time()
    idx = 0
    for bi in range(0, len(items), args.batch):
        batch = items[bi:bi+args.batch]
        imgs = []
        for rid, path, ul in batch:
            try:
                img = tf(Image.open(path).convert('RGB'))
                imgs.append(img)
            except: 
                imgs.append(None)
        valid = [(i,img) for i,img in enumerate(imgs) if img is not None]
        if not valid: continue
        x = torch.stack([v[1] for v in valid]).to(device)
        with torch.no_grad(), torch.amp.autocast('cuda', enabled=(device=="cuda")):
            feats = model(x)
            if isinstance(feats, dict): feats = feats.get('x_norm_clstoken', feats.get('cls', list(feats.values())[0]))
        feats = feats.cpu().numpy().astype(np.float32)
        feats /= (np.linalg.norm(feats, axis=1, keepdims=True) + 1e-8)
        for k, (i, _) in enumerate(valid):
            rid, path, ul = batch[i]
            features[idx] = feats[k]
            labels_list[idx] = ul
            ids_list[idx] = rid
            idx += 1
        if (bi // args.batch) % 50 == 0:
            print(f"  {idx}/{len(items)} ({100*idx/len(items):.1f}%) {idx/max(time.time()-t0,1):.0f} img/s", flush=True)
    features = features[:idx]
    labels_list = labels_list[:idx]
    ids_list = ids_list[:idx]
    
    # Per-class centroid
    from collections import defaultdict
    sums = defaultdict(lambda: np.zeros(emb_dim, dtype=np.float32))
    counts = defaultdict(int)
    for f, l in zip(features, labels_list):
        sums[l] += f
        counts[l] += 1
    centroids = {l: sums[l] / counts[l] for l in sums}
    # Normalize
    centroids = {l: c / (np.linalg.norm(c) + 1e-8) for l, c in centroids.items()}
    
    # Uncertainty = 1 - cos_sim(feature, own_centroid)
    uncertainty = {}
    for f, l, rid in zip(features, labels_list, ids_list):
        c = centroids.get(l)
        if c is None: continue
        sim = float(f @ c)
        uncertainty[rid] = max(0.0, min(1.0, 1.0 - sim))
    
    # Manifest güncelle
    updated = 0
    for d in sorted(SOURCES.iterdir()):
        if not d.is_dir(): continue
        for mf in ['manifest.jsonl', 'manifest_classification.jsonl', 'manifest_detection.jsonl']:
            mp = d / mf
            if not mp.exists(): continue
            records = []
            with open(mp) as f:
                for line in f:
                    r = json.loads(line)
                    u = uncertainty.get(r.get('id'))
                    if u is not None:
                        r['uncertainty_score'] = round(u, 4)
                        r['uncertainty_source'] = f'{model_name}_centroid_v1'
                        updated += 1
                    records.append(r)
            with open(mp, 'w') as f:
                for r in records:
                    f.write(json.dumps(r, ensure_ascii=False) + '\n')
    
    print(f"\nuncertainty_score güncellendi: {updated:,} kayıt")
    
    # Summary
    vals = list(uncertainty.values())
    vals_sorted = sorted(vals)
    n = len(vals_sorted)
    out = {
        "model": model_name,
        "n_scored": n,
        "percentiles": {f"p{p}": round(vals_sorted[int(n*p/100)], 4) for p in [5,25,50,75,95]} if n else {},
        "n_classes_with_centroid": len(centroids),
        "high_uncertainty_threshold": 0.5,
        "n_above_threshold": sum(1 for v in vals if v > 0.5),
    }
    with open(ROOT / "datasets" / "processed" / "uncertainty_dinov3_report.json", 'w') as f:
        json.dump(out, f, indent=2, ensure_ascii=False)
    print(f"\n{json.dumps(out, indent=2)}")

if __name__ == '__main__':
    main()