File size: 4,563 Bytes
28adb60 39d77a4 aa30eea 28adb60 60ba673 f2e3176 46b1f95 28adb60 39d77a4 28adb60 |
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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 |
import os
import sys
import argparse
import numpy as np
import cv2 as cv
from datasets import DATASETS
if "PYTHONPATH" in os.environ:
root_dir = os.environ["PYTHONPATH"]
else:
root_dir = os.path.join("..", "..")
sys.path.append(root_dir)
from models import MODELS
parser = argparse.ArgumentParser("Evaluation with OpenCV on different models in the zoo.")
parser.add_argument("--model", "-m", type=str, required=True, help="model name")
parser.add_argument("--dataset", "-d", type=str, required=True, help="Dataset name")
parser.add_argument("--dataset_root", "-dr", type=str, required=True, help="Root directory of given dataset")
args = parser.parse_args()
models = dict(
mobilenetv1=dict(
name="MobileNetV1",
topic="image_classification",
modelPath=os.path.join(root_dir, "models/image_classification_mobilenet/image_classification_mobilenetv1_2022apr.onnx"),
topK=5),
mobilenetv1_q=dict(
name="MobileNetV1",
topic="image_classification",
modelPath=os.path.join(root_dir, "models/image_classification_mobilenet/image_classification_mobilenetv1_2022apr-int8-quantized.onnx"),
topK=5),
mobilenetv2=dict(
name="MobileNetV2",
topic="image_classification",
modelPath=os.path.join(root_dir, "models/image_classification_mobilenet/image_classification_mobilenetv2_2022apr.onnx"),
topK=5),
mobilenetv2_q=dict(
name="MobileNetV2",
topic="image_classification",
modelPath=os.path.join(root_dir, "models/image_classification_mobilenet/image_classification_mobilenetv2_2022apr-int8-quantized.onnx"),
topK=5),
ppresnet=dict(
name="PPResNet",
topic="image_classification",
modelPath=os.path.join(root_dir, "models/image_classification_ppresnet/image_classification_ppresnet50_2022jan.onnx"),
topK=5),
ppresnet_q=dict(
name="PPResNet",
topic="image_classification",
modelPath=os.path.join(root_dir, "models/image_classification_ppresnet/image_classification_ppresnet50_2022jan-act_int8-wt_int8-quantized.onnx"),
topK=5),
yunet=dict(
name="YuNet",
topic="face_detection",
modelPath=os.path.join(root_dir, "models/face_detection_yunet/face_detection_yunet_2022mar.onnx"),
topK=5000,
confThreshold=0.3,
nmsThreshold=0.45),
yunet_q=dict(
name="YuNet",
topic="face_detection",
modelPath=os.path.join(root_dir, "models/face_detection_yunet/face_detection_yunet_2022mar-act_int8-wt_int8-quantized.onnx"),
topK=5000,
confThreshold=0.3,
nmsThreshold=0.45),
sface=dict(
name="SFace",
topic="face_recognition",
modelPath=os.path.join(root_dir, "models/face_recognition_sface/face_recognition_sface_2021dec.onnx")),
sface_q=dict(
name="SFace",
topic="face_recognition",
modelPath=os.path.join(root_dir, "models/face_recognition_sface/face_recognition_sface_2021dec-act_int8-wt_int8-quantized.onnx")),
crnn=dict(
name="CRNN",
topic="text_recognition",
modelPath=os.path.join(root_dir, "models/text_recognition_crnn/text_recognition_CRNN_EN_2021sep.onnx")),
)
datasets = dict(
imagenet=dict(
name="ImageNet",
topic="image_classification",
size=224),
widerface=dict(
name="WIDERFace",
topic="face_detection"),
lfw=dict(
name="LFW",
topic="face_recognition",
target_size=112),
icdar=dict(
name="ICDAR",
topic="text_recognition"),
iiit5k=dict(
name="IIIT5K",
topic="text_recognition"),
)
def main(args):
# Instantiate model
model_key = args.model.lower()
assert model_key in models
model_name = models[model_key].pop("name")
model_topic = models[model_key].pop("topic")
model_handler, _ = MODELS.get(model_name)
model = model_handler(**models[model_key])
# Instantiate dataset
dataset_key = args.dataset.lower()
assert dataset_key in datasets
dataset_name = datasets[dataset_key].pop("name")
dataset_topic = datasets[dataset_key].pop("topic")
dataset = DATASETS.get(dataset_name)(root=args.dataset_root, **datasets[dataset_key])
# Check if model_topic matches dataset_topic
assert model_topic == dataset_topic
# Run evaluation
dataset.eval(model)
dataset.print_result()
if __name__ == "__main__":
main(args)
|