[Python] BriaAI / Image-Segement(이미지 분리) 누끼따기

 

사용모델 : briai/RMBG-14

https://huggingface.co/briaai/RMBG-1.4

briaai/RMBG-1.4 · Hugging Face

 

 

코드 분석

 

[1] import 

import numpy as np

import torch

from PIL import Image

from skimage import io

import torch.nn.functional as F

#모델 관련 라이브러리

from transformers import AutoModelForImageSegmentation

from torchvision.transforms.functional import normalize

 

#사전 학습된 모델 불러오기

model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-1.4",trust_remote_code=True)

설명
1. 필요한 라이브리러 import
2. 사전 학습된 모델 불러오기

 

 

 

[2] 전처리 후처리 함수 만들기

#데이터 전처리

def preprocess_image(im: np.ndarray, model_input_size: list) -> torch.Tensor:

   

    #1-1. 예외처리

    if len(im.shape) < 3:

        im = im[:, :, np.newaxis]

    #1-2. 차원변경

    im_tensor = torch.tensor(im, dtype=torch.float32).permute(2,0,1)

    #1-3.이미지 입력크기 조정

    im_tensor = F.interpolate(torch.unsqueeze(im_tensor,0), size=model_input_size, mode='bilinear')

    # 1-4 정규화

    image = torch.divide(im_tensor,255.0)

    image = normalize(image,[0.5,0.5,0.5],[1.0,1.0,1.0])

    return image

#데이터 후처리

def postprocess_image(result: torch.Tensor, im_size: list)-> np.ndarray:

    result = torch.squeeze(F.interpolate(result, size=im_size, mode='bilinear') ,0)

    ma = torch.max(result)

    mi = torch.min(result)

    result = (result-mi)/(ma-mi)

    im_array = (result*255).permute(1,2,0).cpu().data.numpy().astype(np.uint8)

    im_array = np.squeeze(im_array)

    return im_array

 

 

[1]데이터 전처리 함수 = preprocess_image

1-1. 흑백인 경우 채널이 2개 -> 3개로 맞춰주기 위해 예외처리
1-2. 차원변경 : (높이,너비,채널) -> (채널 높이 너비)로 변경
        Why? Pytorch 사용 모델의 규격
1-3. 이미지 입력 사이즈 조정   ->  (배치크기,채널,높이,너비)로 수정 

    #torch.unsqueeze(텐서,0) : 텐서의 0번째 차원에 1을 추가하는 것 (배치 크기,채널,높이,너비)

    #bilinear : 보간(보정)법 : 손실을 줄이기 위해 4개의 인접 픽셀에서 값을 이용하여 새로은 픽셀값을 만든다.

1-4. 정규화 
    # 정규화 0-1사이 범주로   
    # 평균값은 0.5 , 표준편차 1.0 고정  


[2] 데이터 후처리 함수 = postprocess_image
2-1. 이미지 출력사이즈 조정 (배치크기 차원 제거)
2-2. 정규화 되돌리기
2-3. numpy 형식 (채널,높이,너비)로 변경
2-4. 필요없는 차원 (길이가 1인 차원) 제거

한마디로 이미지화 시키기 다시 

 

 

[3] 모델 GPU 사용가능하면 옮기기

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model.to(device)

 

 

 

[4] 원하는 사진 ㄱㄱ

# prepare input

image_path = "./누끼테스트3.jpg"

orig_im = io.imread(image_path)

orig_im_size = orig_im.shape[0:2]

image = preprocess_image(orig_im, list(orig_im_size)).to(device)

# inference

result=model(image)

# post process

result_image = postprocess_image(result[0][0], orig_im_size)

# save result

pil_im = Image.fromarray(result_image)

no_bg_image = Image.new("RGBA", pil_im.size, (0,0,0,0))

orig_image = Image.open(image_path)

no_bg_image.paste(orig_image, mask=pil_im)

no_bg_image.save("가나다라마바.png")

 

 

 

참고

사진만 넣고 테스트하고싶은경우~~~~~~~~~~~~~~~~~~~~

https://huggingface.co/spaces/briaai/BRIA-RMBG-1.4

BRIA RMBG 1.4 - a Hugging Face Space by briaai