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 



672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
 
# File 'lib/informers/pipelines.rb', line 672

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