이미지 segmentation 파이썬 예제

💡 Callout

미리 학습된 모델을 기반으로 이미지 segmentation을 수행하여 원본과 model precition 결과를 시각화합니다.

학습 모델은 deeplabv3_resnet101 을 사용합니다.

---

1. conda 환경 설정

제 pc 환경은 gpu가 없었기 때문에 torch를 사용하기 위해서 cpu only로 설치를 진행했습니다.

가상환경은 base가 아닌 미리 생성한 환경에서 진행합니다.

conda install matplotlib
conda install pytorch torchvision torchaudio cpuonly -c pytorch
pip install tensorflow

2. 코드 구성

먼저 아래 코드는 main.py로 이미지가 있어야 할 폴더의 관계는 다음과 같습니다. 

evaluation을 진행할 이미지는 images 폴더 안에 위치해야 합니다. 이미지 이름은 예시처럼 바꾸거나 원하는 이름으로 코드에서 수정하시면 됩니다.

/workspace
  L	/images
       	L test0001.jpg
  L main.py

main.py는 아래와 같습니다.

import torchvision
import torch
import torchvision.transforms as T
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import os

def preprocess_image(image_path):
  # Define the standard transformations
  transform = T.Compose([
      T.Resize((380,520)),
      T.ToTensor(),
      T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
  ])

  # Load the image
  image = Image.open(image_path).convert("RGB")

  # Apply the transformations
  image = transform(image)
  
  return image

def predict(image, model, device):
  # Move the image to the GPU if available
  image = image.to(device)

  # Add a batch dimension
  image = image.unsqueeze(0)

  # Perform the forward pass
  with torch.no_grad():
      output = model(image)['out'][0]

  # Get the predicted class for each pixel
  output_predictions = output.argmax(0)

  return output_predictions.cpu().numpy()

def decode_segmap(image, nc=21):
  label_colors = np.array([(0, 0, 0),  # 0=background
    # Below are the colors for the remaining classes
    (128, 0, 0), (0, 128, 0), (128, 128, 0),
    (0, 0, 128), (128, 0, 128), (0, 128, 128),
    (128, 128, 128), (64, 0, 0), (192, 0, 0),
    (64, 128, 0), (192, 128, 0), (64, 0, 128),
    (192, 0, 128), (64, 128, 128), (192, 128, 128),
    (0, 64, 0), (128, 64, 0), (0, 192, 0),
    (128, 192, 0), (0, 64, 128)])

  r = np.zeros_like(image).astype(np.uint8)
  g = np.zeros_like(image).astype(np.uint8)
  b = np.zeros_like(image).astype(np.uint8)

  for l in range(0, nc):
    idx = image == l
    r[idx] = label_colors[l, 0]
    g[idx] = label_colors[l, 1]
    b[idx] = label_colors[l, 2]

  rgb = np.stack([r, g, b], axis=2)
  return rgb

def visualize_segmentation(image_path, seg_mask):
  # Load the original image
  original_image = Image.open(image_path).convert("RGB")
  original_image = original_image.resize((520, 380))

  # Decode the segmentation mask
  rgb_segmentation = decode_segmap(seg_mask)

  # Plot the original image and the segmentation side by side
  plt.figure(figsize=(10, 5))

  plt.subplot(1, 2, 1)
  plt.title("Original Image")
  plt.imshow(original_image)

  plt.subplot(1, 2, 2)
  plt.title("Segmented Image")
  plt.imshow(rgb_segmentation)

  plt.show()

if __name__ == "__main__":
  
  # Load the pre-trained model
  model = torchvision.models.segmentation.deeplabv3_resnet101(pretrained=True)

  # Set the model to evaluation mode
  model.eval()

  # Check if a GPU is available and if so, move the model to the GPU
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  model.to(device)

  image_path = os.path.join(os.getcwd(), 'images', 'test0001.png')
  # image_files = os.listdir(image_path)
  
  image = preprocess_image(image_path)

  # Predict the segmentation mask
  seg_mask = predict(image, model, device)

  # Visualize the segmentation result
  visualize_segmentation(image_path, seg_mask)

 

핵심이 되는 부분은 미리 학습된 모델을 불러오는 라인입니다.

model = torchvision.models.segmentation.deeplabv3_resnet101(pretrained=True)

이 부분은 원하는 모델로 교체가 가능하기 때문에 다른 모델로도 테스트가 가능합니다.

 

실행하려면 main.py가 있는 곳에서 터미널을 켜고 스크립트에 실행 권한을 주고 다음 명령으로 수행합니다.

chmod +x main.py
python main.py

 

아래 이미지 결과를 보면 segmentation이 잘 된 것을 확인할 수 있습니다.

result

다음에는 instance segmentation으로 고양이들을 분리시켜보겠습니다.