Method: Informers::ImageClassificationPipeline#call

Defined in:

lib/informers/pipelines.rb

#call(images, top_k: 1) ⇒ Object

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

def call(images, top_k: 1)
  is_batched = images.is_a?(Array)
  prepared_images = prepare_images(images)

  pixel_values = @processor.(prepared_images)[:pixel_values]
  output = @model.({pixel_values: pixel_values})

  id2label = @model.config[:id2label]
  to_return = []
  output.logits.each do |batch|
    scores = Utils.get_top_items(Utils.softmax(batch), top_k)

    vals =
      scores.map do |x|
        {
          label: id2label[x[0].to_s],
          score: x[1]
        }
      end
    if top_k == 1
      to_return.push(*vals)
    else
      to_return << vals
    end
  end

  is_batched || top_k == 1 ? to_return : to_return[0]
end