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.



538
539
540
541
542
543
544
545
546
 
# File 'lib/informers/pipelines.rb', line 538

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 



548
549
550
551
552
553
554
555
556
557
558
559
560
561
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
 
# File 'lib/informers/pipelines.rb', line 548

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