MapleStory Finger Point

Development/SAM

SegmentAnyRGBD

吳鍾振 2024. 11. 8. 14:41

SegmentAnyRGBD 코드 분석

https://github.com/Jun-CEN/SegmentAnyRGBD

 

GitHub - Jun-CEN/SegmentAnyRGBD: Segment Any RGBD

Segment Any RGBD. Contribute to Jun-CEN/SegmentAnyRGBD development by creating an account on GitHub.

github.com

 

개요

SegmentAnyRGBD는 RGB 이미지에서 객체 분할(segmentation)과 깊이 추정(depth estimation)을 동시에 수행.
SAM2(Segment Anything Model 2)와 MiDaS(Depth estimation model)를 결합하여 사용합니다.

주요 컴포넌트

1. 초기화

  • SAM2 모델 초기화
    • SAM2ImagePredictor 클래스를 사용하여 segmentation 모델 설정

 

  • MiDaS 모델 초기화
    • Intel의 MiDaS 모델을 PyTorch hub에서 로드
    • DPT_Large 아키텍처 사용

 

  • 임계값 설정 (직접 추가하였음)
    • pred_iou_thresh: 0.86
    • stability_score_thresh: 0.90
    • min_area: 25.0

 

2. 깊이 추정 (estimate_depth)

  • RGB 이미지를 입력받아 깊이 맵 생성
  • 처리 과정:
    1. 이미지 전처리 (MiDaS transform 적용)
    2. 모델을 통한 깊이 예측
    3. 예측 결과를 원본 이미지 크기로 보간
    4. 깊이 맵 정규화

 

3. 세그멘테이션 (segment)

  • 주요 처리 단계:
    1. 이미지 로드 및 전처리
    2. 고해상도 처리를 위한 스케일링 (최소 1024px)
    3. 깊이 추정 수행
    4. SAM2를 통한 마스크 생성
    5. 마스크 후처리 및 정제
    6. 결과를 원본 크기로 복원

 

4. 마스크 정제 (refine_mask_with_depth)

  • 모폴로지 연산을 사용하여 마스크 개선
  • closing 연산으로 작은 구멍을 제거하고 연결성 강화

 

5. 시각화 (visualize_results)

  • matplotlib을 사용하여 결과 시각화
  • 세 가지 이미지 출력:
    1. 원본 RGB 이미지
    2. 추정된 깊이 맵
    3. 세그멘테이션 결과 (컬러 오버레이)

 

메인 실행 로직

  1. 장치 선택 (CUDA/CPU)
  2. SegmentAnyRGBD 인스턴스 생성
  3. 지정된 디렉토리의 이미지들을 순차적으로 처리
  4. 깊이 맵을 컬러 히트맵으로 변환하여 저장

 

사용된 주요 라이브러리

  • OpenCV (cv2)
  • PyTorch
  • NumPy
  • Matplotlib (시각화용)
반응형