본문 바로가기
코드/쉽게 쓰여진 인공지능

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

by 코드퀸 2022. 12. 12.
728x90
반응형

 

 

자율주행차와 
인공지능

 

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

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

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

 

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

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

 

 

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

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

 

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

 

PyTorch

An open source machine learning framework that accelerates the path from research prototyping to production deployment.

pytorch.org

 

 


 

 

오늘의 코드

 

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

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

 

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()
728x90
반응형

댓글