File size: 2,836 Bytes
d9418d2 | 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 | """Verify Stage 1 matches Stage 0 F1 exactly on COCO val 2017."""
import os, sys, json, time
import torch
import numpy as np
from PIL import Image
from pycocotools.coco import COCO
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model import Stage1PersonClassifier
from transformers import AutoModel
COCO_ROOT = '/home/zootest/datasets/coco'
CLASSIFIER_JSON = os.path.join(os.path.dirname(os.path.abspath(__file__)),
'..', 'stage_0', 'classifier.json')
RES = 768
def main():
print('[init] loading Argus + classifier', flush=True)
argus = AutoModel.from_pretrained('/mnt/d/Argus', trust_remote_code=True)
c = json.load(open(CLASSIFIER_JSON))
m = Stage1PersonClassifier(argus, c['pos_dims'], c['neg_dims'], c['threshold']).cuda().eval()
coco = COCO(f'{COCO_ROOT}/annotations/instances_val2017.json')
img_ids = sorted(coco.getImgIds())
labels = [any(a['category_id'] == 1 for a in coco.loadAnns(coco.getAnnIds(imgIds=i, iscrowd=False)))
for i in img_ids]
MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).cuda()
STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).cuda()
scores = torch.zeros(len(img_ids), device='cuda')
t0 = time.time()
for i, img_id in enumerate(img_ids):
info = coco.loadImgs(img_id)[0]
path = f"{COCO_ROOT}/val2017/{info['file_name']}"
img = Image.open(path).convert('RGB').resize((RES, RES), Image.BILINEAR)
arr = np.asarray(img, dtype=np.uint8).copy()
x = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).cuda().float() / 255.0
x = (x - MEAN) / STD
s, _ = m(x)
scores[i] = s[0]
if (i + 1) % 500 == 0:
print(f' {i+1}/{len(img_ids)} {(i+1)/(time.time()-t0):.1f} img/s', flush=True)
y = torch.tensor(labels, dtype=torch.bool, device='cuda')
pred = scores > m.threshold
tp = (pred & y).sum().float()
fp = (pred & ~y).sum().float()
fn = (~pred & y).sum().float()
prec = tp / (tp + fp).clamp(min=1)
rec = tp / (tp + fn).clamp(min=1)
f1 = 2 * prec * rec / (prec + rec).clamp(min=1e-9)
print(f'\n[verify] F1={f1:.4f} P={prec:.4f} R={rec:.4f}', flush=True)
print(f'[parity] Stage 0 baseline = 0.8886; Stage 1 should be identical', flush=True)
out = {
'stage': 1,
'F1': float(f1),
'precision': float(prec),
'recall': float(rec),
'threshold': float(m.threshold.item()),
'stage_0_baseline_F1': 0.8886,
'match': bool(abs(float(f1) - 0.8886) < 1e-3),
}
with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'eval.json'), 'w') as f:
json.dump(out, f, indent=2)
print(f'[done] wrote eval.json match={out["match"]}', flush=True)
if __name__ == '__main__':
main()
|