Class: Informers::ImageClassificationPipeline
- 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, 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 |