Class: Informers::ImageClassificationPipeline

Inherits:
Pipeline
  • Object
show all
Defined in:
lib/informers/pipelines.rb

Instance Method Summary collapse

Methods inherited from Pipeline

#initialize

Constructor Details

This class inherits a constructor from Informers::Pipeline

Instance Method Details

#call(images, top_k: 1) ⇒ Object 



521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
 
# 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