Class: Informers::ImageSegmentationPipeline

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

Instance Method Summary collapse

Constructor Details

#initialize(**options) ⇒ ImageSegmentationPipeline

Returns a new instance of ImageSegmentationPipeline.



552
553
554
555
556
557
558
559
560
 
# File 'lib/informers/pipelines.rb', line 552

def initialize(**options)
  super(**options)

  @subtasks_mapping = {
    "panoptic" => "post_process_panoptic_segmentation",
    "instance" => "post_process_instance_segmentation",
    "semantic" => "post_process_semantic_segmentation"
  }
end

Instance Method Details

#call(images, threshold: 0.5, mask_threshold: 0.5, overlap_mask_area_threshold: 0.8, label_ids_to_fuse: nil, target_sizes: nil, subtask: nil) ⇒ Object 



562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
 
# File 'lib/informers/pipelines.rb', line 562

def call(
  images,
  threshold: 0.5,
  mask_threshold: 0.5,
  overlap_mask_area_threshold: 0.8,
  label_ids_to_fuse: nil,
  target_sizes: nil,
  subtask: nil
)
  is_batched = images.is_a?(Array)

  if is_batched && images.length != 1
    raise Error, "Image segmentation pipeline currently only supports a batch size of 1."
  end

  prepared_images = prepare_images(images)
  image_sizes = prepared_images.map { |x| [x.height, x.width] }

  model_inputs = @processor.(prepared_images).slice(:pixel_values, :pixel_mask)
  output = @model.(model_inputs)

  if !subtask.nil?
    fn = @subtasks_mapping[subtask]
  else
    @subtasks_mapping.each do |task, func|
      if @processor.feature_extractor.respond_to?(func)
        fn = @processor.feature_extractor.method(func)
        subtask = task
        break
      end
    end
  end

  id2label = @model.config[:id2label]

  annotation = []
  if subtask == "panoptic" || subtask == "instance"
    processed = fn.(
      output,
      threshold:,
      mask_threshold:,
      overlap_mask_area_threshold:,
      label_ids_to_fuse:,
      target_sizes: target_sizes || image_sizes, # TODO FIX?
    )[0]

    _segmentation = processed[:segmentation]

    processed[:segments_info].each do |segment|
      annotation << {
        label: id2label[segment[:label_id].to_s],
        score: segment[:score]
        # TODO mask
      }
    end
  elsif subtask == "semantic"
    raise Todo
  else
    raise Error, "Subtask #{subtask} not supported."
  end

  annotation
end