Object Separation (GVTO)

1. 목차

• 환경설정
• 모델 : gvto
• 코드변경
  • 기존 main.py 파일 세그메테이션 적용 수정
• 추론 결과
  • 결과 확인
• 최종 결과

2. 환경 설정

• AI 모델 테스트 환경
  • Ubuntu 22.04(워크스테이션)
  • Anaconda
  • VS Code
  • Python 3.8.18
• 환경설정

git clone <https://github.com/baojudezeze/Generative-Virtual-Try-On>
pip3 install torch torchvision torchaudio --index-url <https://download.pytorch.org/whl/cu121>
pip install -r requirements.txt
pip install pycocotools

디퓨전 모델 사용시
pip install diffusers transformers accelerate scipy safetensors

단일모델 test
pip install segmentation_models_pytorch
pip install diffusers transformers accelerate scipy safetensors
pip install torch torchvision torchaudio
pip install yolov5

3. Generative-Virtual-Try-On

https://github.com/baojudezeze/Generative-Virtual-Try-On

1. 모델 (GVTO)

기존 모델을 사용하여 추론할 때는 pose, mask, face를 분리하고, 분리된 face를 기반으로 모델의 몸을 생성하여 적용한 모델입니다.

4. 코드 변경

main.py

import argparse
import os
import numpy as np
import torch
from PIL import Image
from model.MultimodalAdapter import MultimodalAdapter

# 인자 파싱 함수
def parse_args():
    parser = argparse.ArgumentParser(description="이미지에서 의복 부분을 변경하는 프로그램입니다.")
    parser.add_argument("--device", type=str, default="cuda", choices=['cuda', 'cpu'])
    parser.add_argument("--reference_image_path", type=str, required=True, help="원본 이미지 경로")
    parser.add_argument("--segmentation_image_path", type=str, required=True, help="세그멘테이션 이미지 경로")
    parser.add_argument("--text_prompt", type=str, default="a man wearing a blue jacket", help="생성할 의복의 텍스트 프롬프트")
    parser.add_argument("--inject_cloth_prompt", type=bool, default=False, help="의복 프롬프트를 삽입할지 여부")
    parser.add_argument("--face_scale", type=float, default=0.7, help="얼굴 크기 스케일")
    parser.add_argument("--cloth_scale", type=float, default=0.6, help="의복 크기 스케일")
    parser.add_argument("--num_inference_steps", type=int, default=50, required=True, help="인페인팅 과정에서 사용할 추론 단계 수")
    parser.add_argument("--output_dir", type=str, default="output", help="결과 이미지를 저장할 디렉토리 경로")
    parser.add_argument("--vae_model_path", type=str, default="stabilityai/sd-vae-ft-mse", help="VAE 모델 경로")
    args = parser.parse_args()
    return args

# 마스크 생성 함수
def create_mask(segmentation_image_path):
    # 세그멘테이션 이미지 로드 및 RGBA로 변환
    segmentation_image = Image.open(segmentation_image_path).convert("RGBA")
    
    # 마스크 생성 (빨간색 부분을 마스크로 설정)
    data = np.array(segmentation_image)
    red = [255, 0, 0, 255]
    mask = np.all(data == red, axis=-1)
    
    # 새로운 마스크 이미지 생성
    mask_image = Image.fromarray((mask * 255).astype(np.uint8), mode='L')
    
    return mask_image

# 인페인팅 함수
def inpaint_image(reference_image_path, mask_image, text_prompt, num_inference_steps, args):
    # GVTO 모델 초기화 및 로드
    pipe = MultimodalAdapter(args=args, num_tokens=257)
    
    # 원본 이미지 로드
    original_image = Image.open(reference_image_path).convert("RGB")
    original_array = np.array(original_image)
    
    # 마스크 이미지 로드
    mask_array = np.array(mask_image)

    # 마스크 영역에 대해 인페인팅 수행
    edited_image = pipe.MultimodalPipe(
        seed=66,
        ref_img=original_array,
        face_scale=args.face_scale,
        cloth_scale=args.cloth_scale,
        width=original_array.shape[1],
        height=original_array.shape[0],
        text_prompt=text_prompt,
        num_inference_steps=num_inference_steps
    )[0]
    
    # 편집된 이미지를 배열로 변환
    edited_array = np.array(edited_image)
    
    # 마스크 영역에 대해 편집된 이미지와 원본 이미지를 결합
    composite_image = np.where(mask_array[..., None], edited_array, original_array)
    
    # 결합된 이미지를 PIL 이미지로 변환
    result_image = Image.fromarray(composite_image)
    
    return result_image

# 메인 함수
if __name__ == '__main__':
    # 인자 파싱
    args = parse_args()

    # 출력 디렉토리 생성
    os.makedirs(args.output_dir, exist_okey)

    # 세그멘테이션 이미지에서 마스크 생성
    mask_image = create_mask(args.segmentation_image_path)
    
    # 마스크 이미지 저장 및 표시 (디버깅 용도)
    mask_image_path = os.path.join(args.output_dir, 'mask_image.png')
    mask_image.save(mask_image_path)

    # 인페인팅 수행
    edited_image = inpaint_image(
        reference_image_path=args.reference_image_path,
        mask_image=mask_image,
        text_prompt=args.text_prompt,
        num_inference_steps=args.num_inference_steps,
        args=args
    )
    
    # 결과 이미지 저장
    output_path = os.path.join(args.output_dir, 'output_image_inpainted.png')
    edited_image.save(output_path)

    print("처리가 완료되었습니다. 이미지는 다음 경로에 저장되었습니다:", output_path)

원본 프레임에서 옷 부분을 분리하기 위해 cloth-segmentation 모델을 사용하고, 이 데이터를 기반으로 기존의 세그멘테이션 구조를 변경하여 GVTO 모델을 실증하였습니다.

5. 추론 결과

input Video Frame (GVTO)

output Video Frame (GVTO)

cloth-segmentation 모델을 사용해 분리한 의복 세그메테이션부분에서 상의 부분 Shirt Mask 추출 해서 적용

 

상의 색상 변경

 

상의 세그먼테이션 부분에 옷 이미지 적용

 

Stable Diffusion 기반 생성 적용

6. 최종 결과

색상 변경은 비교적 간단하나, 기존 의류 이미지를 적용하려면 추가적인 학습과 분석이 필요해 보입니다.

Stable Diffusion을 사용하여 생성하고 선택된 영역에 적용했을 때, 버전 문제로 인해 부자연스럽게 보일 수 있지만, 의복 변경이 가능한 것으로 보입니다.

그러나 Prompt를 이용한 생성 방식이므로, 비디오에서 프레임 간 일관성이 유지될지는 확신할 수 없습니다.

GAN 기반 모델을 사용하여 선택된 세그멘테이션 영역에 노이즈를 추가하고 복원하면서 의복 이미지를 생성할 경우에도, 프레임 간 일관성 유지에 대한 예측이 어렵습니다.