Class: Informers::ObjectDetectionPipeline

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, threshold: 0.9, percentage: false) ⇒ Object 



658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
 
# 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