Skip to content

Instantly share code, notes, and snippets.

@ankane
Last active May 10, 2023 17:58
Show Gist options
  • Save ankane/4a9681c8d9b9e814debe9e3ea836529d to your computer and use it in GitHub Desktop.
Save ankane/4a9681c8d9b9e814debe9e3ea836529d to your computer and use it in GitHub Desktop.
require "onnxruntime"
require "mini_magick"
img = MiniMagick::Image.open("bears.jpg")
pixels = img.get_pixels
model = OnnxRuntime::Model.new("model.onnx")
result = model.predict({"inputs" => [pixels]})
p result["num_detections"]
p result["detection_classes"]
coco_labels = {
23 => "bear",
88 => "teddy bear"
}
def draw_box(img, label, box)
width, height = img.dimensions
thickness = 2
top = (box[0] * height).round - thickness
left = (box[1] * width).round - thickness
bottom = (box[2] * height).round + thickness
right = (box[3] * width).round + thickness
# draw box
img.combine_options do |c|
c.draw "rectangle #{left},#{top} #{right},#{bottom}"
c.fill "none"
c.stroke "red"
c.strokewidth thickness
end
# draw text
img.combine_options do |c|
c.draw "text #{left},#{top - 5} \"#{label}\""
c.fill "red"
c.pointsize 18
end
end
result["num_detections"].each_with_index do |n, idx|
n.to_i.times do |i|
label = result["detection_classes"][idx][i].to_i
label = coco_labels[label] || label
box = result["detection_boxes"][idx][i]
draw_box(img, label, box)
end
end
img.write("labeled.jpg")
@ankane
Copy link
Author

ankane commented May 10, 2023

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment