File size: 2,457 Bytes
1387a61 | 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 | import argparse
import os
import cv2 as cv
import numpy as np
here = os.path.dirname(os.path.abspath(__file__))
# TF Object Detection API heads, in the order the graph declares them.
OUTPUTS = ["detection_boxes:0", "detection_scores:0", "detection_classes:0", "num_detections:0"]
def main():
parser = argparse.ArgumentParser(description="Macbeth colour chart detection (ONNX) demo")
parser.add_argument("--model", default=os.path.join(here, "macbeth_chart_detector_2026sep.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("--conf", type=float, default=0.3, help="confidence threshold")
args = parser.parse_args()
img = cv.imread(args.image)
if img is None:
raise SystemExit("could not read image: %s" % args.image)
# The network takes raw uint8 RGB in NHWC layout at the image's own resolution:
# the graph resizes internally (keep-aspect-ratio, min side 768 / max side 1024),
# so do not pre-resize and do not use blobFromImage (that would give NCHW).
blob = cv.cvtColor(img, cv.COLOR_BGR2RGB)[None]
net = cv.dnn.readNetFromONNX(args.model, cv.dnn.ENGINE_OPENCV)
net.setInput(blob)
boxes, scores, _classes, _num = net.forward(OUTPUTS)
boxes = boxes.reshape(-1, 4)
scores = scores.reshape(-1)
# Threshold on the scores rather than on num_detections: that head reports the
# post-NMS slot count (100), not the number of charts actually found. The single
# class is the chart itself, so detection_classes carries no information here.
keep = np.nonzero(scores >= args.conf)[0]
h, w = img.shape[:2]
out = img.copy()
print("macbeth_chart_detector", len(keep), "detections")
for i in keep:
ymin, xmin, ymax, xmax = boxes[i] # normalized, y first
p1 = (int(round(xmin * w)), int(round(ymin * h)))
p2 = (int(round(xmax * w)), int(round(ymax * h)))
cv.rectangle(out, p1, p2, (0, 255, 0), 2)
cv.putText(out, "chart %.2f" % scores[i], (p1[0], max(0, p1[1] - 5)),
cv.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
print(" %.3f x1=%d y1=%d x2=%d y2=%d" % (scores[i], p1[0], p1[1], p2[0], p2[1]))
cv.imwrite(args.output, out)
print("wrote", args.output)
if __name__ == "__main__":
main()
|