Class: Informers::ObjectDetectionPipeline
- Defined in:
- lib/informers/pipelines.rb
Instance Method Summary collapse
Methods inherited from Pipeline
Constructor Details
This class inherits a constructor from Informers::Pipeline
Instance Method Details
#call(images, threshold: 0.9, percentage: false) ⇒ Object
672 673 674 675 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693 694 695 696 697 698 699 700 701 702 703 |
# File 'lib/informers/pipelines.rb', line 672 def call(images, threshold: 0.9, percentage: false) is_batched = images.is_a?(Array) if is_batched && images.length != 1 raise Error, "Object detection pipeline currently only supports a batch size of 1." end prepared_images = prepare_images(images) image_sizes = percentage ? nil : prepared_images.map { |x| [x.height, x.width] } model_inputs = @processor.(prepared_images).slice(:pixel_values, :pixel_mask) output = @model.(model_inputs) processed = @processor.feature_extractor.post_process_object_detection(output, threshold, image_sizes) # Add labels id2label = @model.config[:id2label] # Format output result = processed.map do |batch| batch[:boxes].map.with_index do |box, i| { label: id2label[batch[:classes][i].to_s], score: batch[:scores][i], box: get_bounding_box(box, !percentage) } end.sort_by { |v| -v[:score] } end is_batched ? result : result[0] end |