KHU_4-1_JP

생성형 AI를 이용한 시각적 질감 기반 하이트맵 생성 연구

본 프로젝트는 생성형 AI 모델을 사용하여 2D 질감 이미지로부터 3D 표면의 굴곡 정보를 담은 하이트맵(Heightmap)을 생성하는 연구입니다. Gelsight 센서로 촬영된 시각적 질감 정보를 입력받아, 해당 표면의 미세한 지오메트리를 추론하여 현실적인 촉각 피드백의 기반 데이터를 생성하는 것을 목표로 합니다.

이 프로젝트는 경희대학교 컴퓨터공학과에서 진행된 연구 “생성형 AI를 이용한 시각적 질감 기반 하이트맵 생성 연구”의 공식 구현체입니다.

1. 개요

디지털 콘텐츠의 몰입감을 높이는 데 현실적인 촉각 피드백은 매우 중요합니다. 본 연구에서는 U-Net 기반의 이미지 대 이미지(Image-to-Image) 변환 모델을 사용하여, Stanford-Gelsight 데이터셋의 RGB 질감 이미지로부터 직접 하이트맵을 생성하는 파이프라인을 구축했습니다.

2. 시스템 아키텍처

본 모델은 segmentation_models_pytorch 라이브러리를 기반으로 한 U-Net 아키텍처를 사용합니다.

  1. 인코더 (Encoder): ImageNet으로 사전 학습된 EfficientNet-B7 모델을 사용하여 입력 이미지로부터 풍부한 시각적 특징을 추출합니다.
  2. 디코더 (Decoder): 인코더에서 압축된 특징을 점진적으로 업샘플링하여 원본 크기의 하이트맵을 생성합니다.
  3. 스킵 커넥션 (Skip Connection): 인코더의 각 레벨에서 추출된 저수준 특징(Low-level features)을 디코더의 해당 레벨로 직접 전달하여, 경계선 등 세밀한 디테일이 소실되는 것을 방지합니다.

3. 설치 방법

본 프로젝트는 Conda 가상환경에서의 실행을 권장합니다.

  1. 저장소 클론 
    git clone https://github.com/your-username/your-repository-name.git
cd your-repository-name

  2. Conda 가상환경 생성 및 활성화 environment.yml 파일을 사용하여 필요한 모든 패키지를 설치합니다.

    conda env create -f environment.yml
conda activate JP_3rdTry
  3. 주요 의존성 패키지
    • python=3.12
    • pytorch & torchvision
    • opencv-python
    • scikit-image
    • segmentation-models-pytorch
    • lpips

4. 사용 방법

프로젝트의 주요 기능은 다음 스크립트를 통해 실행할 수 있습니다.

데이터셋 준비

  1. Stanford-Gelsight 데이터셋을 다운로드하여 IOFiles/Input과 유사한 구조로配置합니다.
  2. dataset.py의 TextureHeightmapDataset 클래스는 이미지 로딩, 전처리(마커 제거, 정규화 등) 파이프라인을 포함하고 있습니다.

모델 학습

run_train.sh 셸 스크립트를 실행하거나, re_train.py를 직접 실행하여 모델 학습을 시작할 수 있습니다.

# 스크립트 사용 (권장)
bash re_run_train.sh

# 직접 실행
python re_train.py --encoder_name efficientnet-b7 --batch_size 32 --lr 1e-4 --lambda_lpips 0.5 --epochs 1000

모델 평가

re_evaluate.py 스크립트를 사용하여 학습된 모델의 성능을 정량적으로 평가합니다.

python re_evaluate.py --model_path /path/to/your/model.pth --data_root /path/to/test/dataset

하이트맵 생성 (추론)

re_predict.py 스크립트를 사용하여 새로운 질감 이미지로부터 하이트맵을 생성합니다.

python re_predict.py --model_path /path/to/your/model.pth --input_image /path/to/input/image.png --output_path /path/to/output/heightmap.png

5. 결과 및 분석

결론적으로, 제안된 모델은 질감의 전반적인 느낌은 잘 재현하지만, 디테일과 정확성 측면에서는 개선의 여지가 있습니다.

5.1. 시각적 결과 예시

성공적인 사례 (Successful Cases)

전반적인 구조와 질감의 통계적 특성을 잘 재현한 경우입니다. (예: PlasticMesh2', 'ArtificialGrass)

입력 (Input) 원본 (Ground Truth) 생성 (Generated)
Success case example 1
Success case example 2


한계가 보이는 사례 (Cases with Limitations)

세밀한 패턴이나 날카로운 경계선을 재현하는 데 어려움을 겪는 경우입니다. (예: PaintedWood)

입력 (Input) 원본 (Ground Truth) 생성 (Generated)
Limitation case example

6. 향후 연구 방향

  1. 손실 함수 고도화: GLCM Contrast를 직접 제어하는 손실 항을 추가하여 특징 과장 문제를 완화하고, Gradient 기반 손실을 통해 엣지 표현력을 강화.
  2. 아키텍처 탐색: U-Net 외에 GAN(Generative Adversarial Network)과 같은 다른 생성 모델 아키텍처를 도입하여 엣지 생성 능력 개선.
  3. 데이터셋 확장: 자체 Gelsight 장비를 활용하여 더 다양하고 고품질의 데이터를 수집하여 모델의 일반화 성능 향상.

7. 참고 문헌