| """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() |
|
|
