Skip to content

Instantly share code, notes, and snippets.

@ankane ankane/detect_objects.rb
Last active Aug 29, 2019

Embed
What would you like to do?
require "onnxruntime"
require "mini_magick"
img = MiniMagick::Image.open("bears.jpg")
pixels = img.get_pixels
model = OnnxRuntime::Model.new("model.onnx")
result = model.predict("image_tensor:0" => [pixels])
p result["num_detections:0"]
p result["detection_classes:0"]
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:0"].each_with_index do |n, idx|
n.to_i.times do |i|
label = result["detection_classes:0"][idx][i].to_i
label = coco_labels[label] || label
box = result["detection_boxes:0"][idx][i]
draw_box(img, label, box)
end
end
img.write("labeled.jpg")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.