Pytorch로 Face Segmentation 해보기 간단 예제

안녕하세요. 이번 글에서는

파이토치로 얼굴 이미지에 대하여 부위별 세그멘테이션을 진행하는 

아주 간단한 예제 코드를 다루어보도록 하겠습니다.

 

먼저, 필요한 모듈들을 설치해줍니다.

pytorch의 경우에는 반드시 CUDA의 버전과 호환되도록 맞추어서 설치해주셔야 합니다.

$pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124  # CUDA 12.4용 예시
$pip install transformers pillow matplotlib requests

 

모듈 설치가 완료되었다면 아래의 코드 스니펫으로 Face Segmentation을

진행해보실 수 있습니다. 초기 코드 실행 시에 한해서 모델 다운로드가 진행됩니다.

import torch, requests
from PIL import Image
from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
import matplotlib.pyplot as plt
import numpy as np

# 디바이스 설정(cuda 권장)
device = "cuda" if torch.cuda.is_available() else "cpu"

# 모델 다운로드(출처: https://huggingface.co/jonathandinu/face-parsing)
processor = SegformerImageProcessor.from_pretrained("jonathandinu/face-parsing")
model = SegformerForSemanticSegmentation.from_pretrained("jonathandinu/face-parsing").to(device)

# 예시 이미지 다운로드(unsplash 사이트에 있는 얼굴 이미지 예시)
url = "https://images.unsplash.com/photo-1539571696357-5a69c17a67c6?q=80&w=687&auto=format&fit=crop&ixlib=rb-4.1.0&ixid=M3wxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8fA%3D%3D"
image = Image.open(requests.get(url, stream=True).raw)

# 이미지 전처리 및 추론 진행
inputs = processor(images=image, return_tensors="pt").to(device)
outputs = model(**inputs)
logits = torch.nn.functional.interpolate(outputs.logits, size=image.size[::-1],
                                         mode="bilinear", align_corners=False)
mask = logits.argmax(dim=1)[0].cpu().numpy()

# 추론된 각 마스크 부위마다 특정 컬러로 맵핑
palette = np.random.randint(0, 255, (19, 3), dtype=np.uint8) # 랜덤하게 색상 선택
color_mask = palette[mask]
overlay = Image.blend(image.convert("RGBA"),
                      Image.fromarray(color_mask, mode="RGB").convert("RGBA"),
                      alpha=0.5)

# 시각화
plt.subplot(1,2,1); plt.imshow(image); plt.axis("off"); plt.title("Original")
plt.subplot(1,2,2); plt.imshow(overlay); plt.axis("off"); plt.title("Segmentation")
plt.tight_layout(); plt.show()

 

각 부위에 대한 정보 / 더 상세한 예제 등이 필요하시다면 아래의 링크 정보를 참고하세요.

jonathandinu/face-parsing · Hugging Face

 

이 글이 파이토치를 활용한 Face Segmentation 과정에 도움이 되셨기를 기원합니다.

감사합니다.