| """Evaluate Argus's heads on their benchmarks. |
| |
| python evaluate.py imagenet --root VAL_DIR # VAL_DIR/<wnid>/*.JPEG |
| python evaluate.py ade20k --root ADEChallengeData2016 |
| python evaluate.py nyu --root NYU_TEST_DIR # rgb/<id>.* and depth/<id>.png |
| python evaluate.py coco --images val2017 --annotations instances_val2017.json |
| |
| Every command takes --model (default phanerozoic/argus), --device, --batch |
| and --limit N to evaluate the first N items only. |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import sys |
|
|
| import numpy as np |
| import torch |
| from PIL import Image |
| from transformers import AutoModel |
|
|
| IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".JPEG", ".JPG", ".PNG") |
|
|
|
|
| def _batches(items, size): |
| for i in range(0, len(items), size): |
| yield items[i:i + size] |
|
|
|
|
| def _load(path): |
| return Image.open(path).convert("RGB") |
|
|
|
|
| def evaluate_imagenet(model, args): |
| """Top-1 and top-5 accuracy with classify(), on an ImageNet validation |
| folder of <wnid>/<image> files.""" |
| wnid_to_idx = {w: i for i, w in enumerate(model.class_ids)} |
| items = sorted( |
| (os.path.join(args.root, w, f), wnid_to_idx[w]) |
| for w in os.listdir(args.root) if w in wnid_to_idx |
| for f in os.listdir(os.path.join(args.root, w)) if f.endswith(IMAGE_EXTS) |
| )[: args.limit] |
| top1 = top5 = 0 |
| for chunk in _batches(items, args.batch): |
| results = model.classify([_load(p) for p, _ in chunk], top_k=5) |
| for (_, label), result in zip(chunk, results): |
| predicted = [wnid_to_idx[r["class_id"]] for r in result] |
| top1 += predicted[0] == label |
| top5 += label in predicted |
| return {"images": len(items), "top1": top1 / len(items), "top5": top5 / len(items)} |
|
|
|
|
| def evaluate_ade20k(model, args): |
| """Mean IoU over the 150 classes with segment(), on the ADEChallengeData2016 |
| validation split, whose annotations store class k as k + 1 and 0 as |
| unlabeled.""" |
| image_dir = os.path.join(args.root, "images", "validation") |
| label_dir = os.path.join(args.root, "annotations", "validation") |
| names = sorted(f for f in os.listdir(image_dir) if f.endswith(IMAGE_EXTS))[: args.limit] |
| n = model.config.num_seg_classes |
| intersection = np.zeros(n, dtype=np.int64) |
| union = np.zeros(n, dtype=np.int64) |
| for chunk in _batches(names, args.batch): |
| preds = model.segment([_load(os.path.join(image_dir, f)) for f in chunk]) |
| for f, pred in zip(chunk, preds): |
| label = np.asarray(Image.open(os.path.join(label_dir, os.path.splitext(f)[0] + ".png")), dtype=np.int64) - 1 |
| pred = pred.cpu().numpy() |
| valid = label >= 0 |
| p, t = pred[valid], label[valid] |
| intersection += np.bincount(t[p == t], minlength=n)[:n] |
| union += np.bincount(p, minlength=n)[:n] + np.bincount(t, minlength=n)[:n] - np.bincount(t[p == t], minlength=n)[:n] |
| present = union > 0 |
| return {"images": len(names), "miou": float((intersection[present] / union[present]).mean())} |
|
|
|
|
| def evaluate_nyu(model, args): |
| """Depth error with depth(), on NYU Depth V2 test images and 16-bit depth |
| PNGs in millimetres, inside the Eigen crop and up to 10 m.""" |
| rgb_dir, depth_dir = os.path.join(args.root, "rgb"), os.path.join(args.root, "depth") |
| stems = sorted(os.path.splitext(f)[0] for f in os.listdir(depth_dir) if f.endswith(".png"))[: args.limit] |
| rgb_files = {os.path.splitext(f)[0]: f for f in os.listdir(rgb_dir)} |
| abs_rel, sq_rmse, delta1, count = 0.0, 0.0, 0.0, 0 |
| for chunk in _batches(stems, args.batch): |
| preds = model.depth([_load(os.path.join(rgb_dir, rgb_files[s])) for s in chunk]) |
| for s, pred in zip(chunk, preds): |
| gt = np.asarray(Image.open(os.path.join(depth_dir, s + ".png")), dtype=np.float64) / 1000.0 |
| pred = pred.cpu().numpy().astype(np.float64) |
| crop = np.zeros_like(gt, dtype=bool) |
| h, w = gt.shape |
| crop[int(0.4081 * h):int(0.9919 * h), int(0.0359 * w):int(0.9640 * w)] = True |
| valid = crop & (gt > 1e-3) & (gt <= 10.0) |
| g, p = gt[valid], np.clip(pred[valid], 1e-3, 10.0) |
| abs_rel += np.mean(np.abs(p - g) / g) |
| sq_rmse += np.sqrt(np.mean((p - g) ** 2)) |
| delta1 += np.mean(np.maximum(p / g, g / p) < 1.25) |
| count += 1 |
| return {"images": count, "abs_rel": abs_rel / count, "rmse": sq_rmse / count, "delta1": delta1 / count} |
|
|
|
|
| def evaluate_coco(model, args): |
| """Box mAP with detect(), scored by pycocotools on COCO val2017.""" |
| from pycocotools.coco import COCO |
| from pycocotools.cocoeval import COCOeval |
|
|
| gt = COCO(args.annotations) |
| name_to_cat = {c["name"]: c["id"] for c in gt.loadCats(gt.getCatIds())} |
| present = set(os.listdir(args.images)) |
| image_ids = [i for i in sorted(gt.getImgIds()) if gt.imgs[i]["file_name"] in present][: args.limit] |
| detections = [] |
| for chunk in _batches(image_ids, args.batch): |
| results = model.detect([_load(os.path.join(args.images, gt.imgs[i]["file_name"])) for i in chunk]) |
| for image_id, result in zip(chunk, results): |
| for d in result: |
| x1, y1, x2, y2 = d["box"] |
| detections.append({ |
| "image_id": image_id, |
| "category_id": name_to_cat[d["class_name"]], |
| "bbox": [x1, y1, x2 - x1, y2 - y1], |
| "score": d["score"], |
| }) |
| evaluator = COCOeval(gt, gt.loadRes(detections), "bbox") |
| evaluator.params.imgIds = image_ids |
| evaluator.evaluate() |
| evaluator.accumulate() |
| evaluator.summarize() |
| s = evaluator.stats |
| return {"images": len(image_ids), "map": s[0], "map50": s[1], "map75": s[2], "map_small": s[3], "map_medium": s[4], "map_large": s[5]} |
|
|
|
|
| TASKS = {"imagenet": evaluate_imagenet, "ade20k": evaluate_ade20k, "nyu": evaluate_nyu, "coco": evaluate_coco} |
|
|
|
|
| def main(argv=None): |
| parser = argparse.ArgumentParser(description="Evaluate Argus's heads on their benchmarks.") |
| parser.add_argument("task", choices=sorted(TASKS)) |
| parser.add_argument("--root", help="dataset root for imagenet, ade20k and nyu") |
| parser.add_argument("--images", help="COCO val2017 image directory") |
| parser.add_argument("--annotations", help="COCO instances_val2017.json") |
| parser.add_argument("--model", default="phanerozoic/argus") |
| parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| parser.add_argument("--batch", type=int, default=8) |
| parser.add_argument("--limit", type=int, default=None) |
| args = parser.parse_args(argv) |
|
|
| model = AutoModel.from_pretrained(args.model, trust_remote_code=True).to(args.device).eval() |
| result = TASKS[args.task](model, args) |
| print(json.dumps({"task": args.task, **result}, indent=2)) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
