isntance segmentation 추출 작업

Jm-baek·2021년 12월 28일

Detectron2

목록 보기

3/4

이미지를 무단으로 사용하지마세요! 신고합니다.

Detectron2에서 원하는 객체를 추출하는 방법(나의 뇌피셜)은 두 가지 방법이 있다.

1번 방법. model을 학습한 후, 직접 범위를 가져와서 작업한다.

예를 들면 아래코드와 같다.

mask_array = outputs["instances"].pred_masks.cpu().numpy()
num_instances = mask_array.shape[0]
mask_array = np.moveaxis(mask_array, 0, -1)

mask_array_instance = []
output = np.zeros_like(image) #black

#print('output',output)
for i in range(num_instances):
    mask_array_instance.append(mask_array[:, :, i:(i+1)])
    # print(mask_array_instance)
    output = np.where(mask_array_instance[i] == True, 255, output)

cv2_imshow(output)

1번 방법 사용

1번 방법은 모르겠다.
seg의 결과값 구성과 rgb에서 rgba까지 사용하던데 모르겠다.^^

하지만 1번 방법처럼 작업을 해야하고 나중에는 1번 방법에 대해 다시 공부하겠지만 지금 급한 불을 끄기위해 편법(?)으로 2번 방법을 진행했다.

2번 방법. Detectron2의 내부 코드를 수정해버렸다.

이미 그려내는 코드가 있기 때문에 내가 원하는 것만 그리도록 수정을 했다.(완벽하지 않다.)

detectron2의 utils 에 가면 visualizer.py 가 있다.
Visualizer 클래스 안에서 draw_panoptic_seg를 수정하면된다.

최종 이미지 그려내는 코드

Visualzier의 클래스가 v 인스턴스를 생성하고 v는 draw_panoptic_seg 메소드를 사용하여 최종 out 변수를 생성하고 있다.
draw_panoptic_seg를 들어가보자.

outputs =  predictor(image)			# 모델 결과 이미지
panoptic_seg, segments_info = outputs["panoptic_seg"]

# class 인스턴스
v = Visualizer(image[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=0.5)
    
# 인스턴스 메소드 사용
out = v.draw_panoptic_seg(panoptic_seg.to("cpu"), segments_info, alpha = 1)

draw_panoptic_seg 함수 수정 부분

위 코드에서 보면 v.draw_panoptic_sge가 그림을 그리는 것을 볼 수 있다.
중요한 부분이 "instance"를 그려내는 코드가 존재하지만 나는 이 부분을 삭제해버렸다.
배경만 없어지면 괜찮기 때문이다.

def draw_panoptic_seg(self, panoptic_seg, segments_info, area_threshold=None, alpha=0.7): 함수 수정 부분

  pred = _PanopticPrediction(panoptic_seg, segments_info, self.metadata)

  if self._instance_mode == ColorMode.IMAGE_BW:
      self.output.reset_image(self._create_grayscale_image(pred.non_empty_mask()))

  # draw mask for all semantic segments first i.e. "stuff"
  for mask, sinfo in pred.semantic_masks():
      # 내가 수정한 부분
      # 21번 category id의 배경은 못 지우게 예외시켰다.
      if sinfo["category_id"] == 21:    
          continue
      category_idx = sinfo["category_id"]
      try:
          # 내가 수정한 부분
          # 필요없는 배경들은 전부 검은색으로 변경
          self.metadata.stuff_colors[category_idx] = [0, 0, 0]  
          mask_color = [x / 255 for x in self.metadata.stuff_colors[category_idx]]
      except AttributeError:
          mask_color = None

          text = self.metadata.stuff_classes[category_idx]
          self.draw_binary_mask(
                  mask,
                  color=mask_color,
                  edge_color=None,	#  None
                  text=None,		#  None	
                  alpha=alpha,
                  area_threshold=None,  #  None
              )