added lapsrn an macbeth onnx (#16)

1387a61 about 1 month ago

2.46 kB

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

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