Method: Informers::ObjectDetectionPipeline#call

Defined in:

lib/informers/pipelines.rb

#call(images, threshold: 0.9, percentage: false) ⇒ Object

# File 'lib/informers/pipelines.rb', line 658

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