[시맨틱 세그멘테이션] 왕초보도 가능한 이미지 영역 분리하기 (python, pytorch)

 

 

자율주행차와 
인공지능

 

자율주행차는 종합예술,미래  기술의 집약체라는 말을 종종하곤 합니다.

요즘 나오는 자율주행차를 보고 있으면 언제 이렇게 기술이 발전했나 싶을정도로 기술의 발전이 놀라운데요,

오늘은 그 일부인 시맨틱 세그멘테이션 기술을 소개합니다.

 

시맨틱 세그멘테이션은 말 그대로, 의미론적으로 나누는 것입니다. 

사람은 사람 영역으로, 길은 길 영역으로 나누는 기술입니다.

 

 

코드는 아래의 pytorch 공식 홈페이지에서 가져왔으며,

아래 링크를 들어가면, 환경설치 없이도 간단하게 코드를 테스트해볼 수  있습니다.

 

https://pytorch.org/hub/pytorch_vision_deeplabv3_resnet101/

PyTorch

 

 

---

 

 

오늘의 코드

 

코드의 설명은 주석을 참고해주세요.

짧진 않지만, 어렵지 않은 코드이니 쭉 보고 각 줄의 의미를 이해하면 좋을 듯 싶습니다. 

 

import torch
from PIL import Image
from torchvision import transforms
import matplotlib.pyplot as plt
torch.hub._validate_not_a_forked_repo=lambda a,b,c: True

# 입력 이미지 불러오기
input_image = Image.open('load.png')
input_image = input_image.convert("RGB")

# 입력 이미지를 딥러닝 모델에 넣을 수 있는 형태로 전처리 수행
preprocess = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])
input_tensor = preprocess(input_image)
input_batch = input_tensor.unsqueeze(0)

#딥러닝 불러오기
model = torch.hub.load('pytorch/vision:v0.10.0', 'deeplabv3_resnet50', pretrained=True)
model.eval()

#딥러닝에 입력 이미지 넣고 예측결과 가져오기
output = model(input_batch)['out'][0]
output_predictions = output.argmax(0)

# 결과를 칼라로 표현하도록 색상표 만들기
palette = torch.tensor([2 ** 25 - 1, 2 ** 15 - 1, 2 ** 21 - 1])
colors = torch.as_tensor([i for i in range(21)])[:, None] * palette
colors = (colors % 255).numpy().astype("uint8")

# 결과 이미지 생성하기
r = Image.fromarray(output_predictions.byte().cpu().numpy()).resize(input_image.size)
r.putpalette(colors)

# 결과 이미지 표출하기
plt.imshow(r)
plt.show()