Image Mask and Polygon

사건의 발단 Stay...!

인공지능 경진대회에 나갔는데 생각보다 점수의 상승폭이 적어서 남들과는 다른 거를 해보기로 했다.
아래의 이미지를 Image Segmentation 모델에 넣어 학습을 시켜보기 위해 폴리곤으로 따내려고 한다.
계획은 cv2의 inRange를 사용하여 특정 색상 영역을 추출하여 분리하고,
폴리곤으로 추출해내려고 한다.

누끼? 따기

hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
lower_black = np.array([0,0,0])
upper_black = np.array([0,0,0])
mask = cv2.inRange(hsv, lower_black, upper_black)
cv2_imshow(mask)

위 코드를 실행하면 아래 이미지가 나온다.
근데 생각처럼 깔끔하게 따지지가 않았다. 중간중간에 흰색이 들어가 있어서 하나의 mask가 되지 않는다.

문제점 1: 사이사이 점들이 있다.

li = []
for line in mask:
  li = line
  res = [i for i, val in enumerate(li) if val == 0]
  st, end = res[0], res[-1]
  for i in range(len(line)):
    if line[i] == 255:
      line[i] = 1
  for i in range(st, end):
    line[i] = 0
mask

위 코드로 검은 점 사이에 있는 흰색 공간들을 채워주었다.
그리고 255를 1로 바꾸어주었다.
아래 그림은 1로 바꾸어주기 이전에 mask를 출력한 결과이다.

문제점 2: 폴리곤을 못 딴다...

하지만 여기서도 문제가 있다.
이미지를 잘 보면 고기가 잘려있는 이미지가 있는데,
그런 경우 폴리라인을 이상하게 딴다... 정신나갈뻔했다.

해결책: 테두리 주기

그래서 낸 아이디어가 위아래의 픽셀값에 4픽셀 정도 강제로 테두리를 주어
폴리곤으로 따내는 것이었다.

def format_mask(mask):
  n_mask = []
  for line_idx in range(len(mask)):
    line_list = list(mask[line_idx])
    for ele_idx in range(len(line_list)):
      if line_list[ele_idx] == 1:
        line_list[ele_idx] = 0
      else:
        line_list[ele_idx] = 1

      if line_idx < 4:
        line_list[ele_idx] = 0
      elif line_idx > len(mask) - 5:
        line_list[ele_idx] = 0
      if ele_idx < 4 or ele_idx > len(line_list) -5 :
        line_list[ele_idx] = 0
    n_mask.append(line_list)
  return n_mask

위의 코드를 실행시키면 상하좌우 4픽셀씩 0으로 고정시킨다.
또한 기존의 0값은 1로 변환시키고 1은 0으로 변환시킨다.

폴리곤 변환

아래는 폴리곤으로 변환시키는 코드이다.

import skimage.measure as measure

def close_contour(contour):
  if not np.array_equal(contour[0], contour[-1]):
    contour = np.vstack((contour, contour[0]))
  return contour

def binary_mask_to_polygon(binary_mask, tolerance=0):
  polygons = []
  # pad mask to close contours of shapes which start and end at an edge
  padded_binary_mask = np.pad(binary_mask, pad_width=1, mode='constant', constant_values=0)
  contours = measure.find_contours(padded_binary_mask, 0.5)
  contours = np.subtract(contours, 1)
  for contour in contours:
      contour = close_contour(contour)
      contour = measure.approximate_polygon(contour, tolerance)
      if len(contour) < 3: 
          continue
      contour = np.flip(contour, axis=1)
      segmentation = contour.ravel().tolist()
      # after padding and subtracting 1 we may get -0.5 points in our segmentation
      segmentation = [0 if i < 0 else i for i in segmentation]
      polygons.append(segmentation)

  return polygons
  
poly = binary_mask_to_polygon(conv_mask)
len(poly)

px = [a for a in poly[0][0::2]]
py = [a for a in poly[0][1::2]]
poly2 = [[x, y] for x, y in zip(px, py)]

전체 코드

최종 코드는 이렇게 된다.

import os
import numpy as np
import json
from detectron2.structures import BoxMode
import cv2
from tqdm import tqdm
import joblib

def format_mask(mask):
  n_mask = []
  for line_idx in range(len(mask)):
    line_list = list(mask[line_idx])
    for ele_idx in range(len(line_list)):
      if line_list[ele_idx] == 1:
        line_list[ele_idx] = 0
      else:
        line_list[ele_idx] = 1

      if line_idx < 4:
        line_list[ele_idx] = 0
      elif line_idx > len(mask) - 5:
        line_list[ele_idx] = 0
      if ele_idx < 4 or ele_idx > len(line_list) -5 :
        line_list[ele_idx] = 0
    n_mask.append(line_list)
  return n_mask

import skimage.measure as measure

def close_contour(contour):
  if not np.array_equal(contour[0], contour[-1]):
    contour = np.vstack((contour, contour[0]))
  return contour

def binary_mask_to_polygon(binary_mask, tolerance=0):
  polygons = []
  # pad mask to close contours of shapes which start and end at an edge
  padded_binary_mask = np.pad(binary_mask, pad_width=1, mode='constant', constant_values=0)
  contours = measure.find_contours(padded_binary_mask, 0.5)
  contours = np.subtract(contours, 1)
  for contour in contours:
      contour = close_contour(contour)
      contour = measure.approximate_polygon(contour, tolerance)
      if len(contour) < 3: 
          continue
      contour = np.flip(contour, axis=1)
      segmentation = contour.ravel().tolist()
      # after padding and subtracting 1 we may get -0.5 points in our segmentation
      segmentation = [0 if i < 0 else i for i in segmentation]
      polygons.append(segmentation)

  return polygons


def get_cow_dicts(directory):
    classes = ['1++', '1+', '1', '2', '3']
    dataset_dicts = []
    for idx, filename in enumerate(tqdm([file for file in os.listdir(directory) if file.endswith('.jpg')])):
        img_file = os.path.join(directory, filename)
        img = cv2.imread(img_file, cv2.IMREAD_COLOR)
        height, width, c = img.shape
        tclass = filename.split('_')[1]
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        lower_black_t = np.array([0,0,0])
        upper_black_t = np.array([20,20,20])
        mask1 = cv2.inRange(hsv, lower_black_t, upper_black_t)

        li = []
        for line in mask1:
          li = line
          res = [i for i, val in enumerate(li) if val == 0]
          st, end = res[0], res[-1]
          for i in range(len(line)):
            if line[i] == 255:
              line[i] = 1
          for i in range(st, end):
            line[i] = 0

        conv_mask = np.array(format_mask(mask1))
        poly = binary_mask_to_polygon(conv_mask)

        px = [a for a in poly[0][0::2]]
        py = [a for a in poly[0][1::2]]
        poly2 = [[x, y] for x, y in zip(px, py)]
        poly2 = [p for x in poly2 for p in x]

        record = {}
        
        record["file_name"] = img_file
        record["height"] = height
        record["width"] = width
        record["image_id"] = idx
      
        objs = []
        
        obj = {
            "bbox": [np.min(px), np.min(py), np.max(px), np.max(py)],
            "bbox_mode": BoxMode.XYXY_ABS,
            "segmentation": [poly2],
            "category_id": classes.index(tclass),
            "iscrowd": 0
        }

        objs.append(obj)
        record["annotations"] = objs
        dataset_dicts.append(record)
    return dataset_dicts

joblib.dump(get_cow_dicts('/content/drive/MyDrive/KoreanBeef/data/00_source/images'), "/content/drive/MyDrive/KoreanBeef/data/pickledata.pkl")

결과

느낀점

이거 한다고 별 짓을 다했던거 같은데
결과는 안좋더라.
역시 분류는 분류로 돌리는게 맞더라.

오늘의 교훈
사람들이 안하는데는 이유가 있다.
그래도 이런 삽질을 하면서 생각보다 많은 것을 배울 수 있었던 좋은 경험이 된 것 같다.