File size: 1,502 Bytes
e9eb33d | 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 | import argparse
import os
import cv2 as cv
import numpy as np
here = os.path.dirname(os.path.abspath(__file__))
def main():
parser = argparse.ArgumentParser(description="TensorFlow Inception (ONNX) image classification demo")
parser.add_argument("--model", default=os.path.join(here, "tensorflow_inception_graph_2026jul.onnx"))
parser.add_argument("--image", default=os.path.join(here, "example_outputs", "input_image.png"))
parser.add_argument("--output", default=os.path.join(here, "example_outputs", "output_image.png"))
parser.add_argument("--labels", help="optional ImageNet label file, one class name per line")
args = parser.parse_args()
img = cv.imread(args.image)
if img is None:
raise SystemExit("could not read image: %s" % args.image)
rgb = cv.resize(cv.cvtColor(img, cv.COLOR_BGR2RGB), (224, 224)).astype(np.float32)
net = cv.dnn.readNetFromONNX(args.model)
net.setInput(rgb[None])
scores = net.forward().ravel()
top = int(np.argmax(scores))
conf = float(scores[top])
label = str(top)
if args.labels:
names = open(args.labels).read().splitlines()
if top < len(names):
label = names[top]
print("class", top, label, "confidence", round(conf, 4))
out = img.copy()
cv.putText(out, "%s (%.2f)" % (label, conf), (10, 30), cv.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 0), 2)
cv.imwrite(args.output, out)
print("wrote", args.output)
if __name__ == "__main__":
main()
|