SegmentAnyRGBD 코드 분석
https://github.com/Jun-CEN/SegmentAnyRGBD
개요
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 이미지를 입력받아 깊이 맵 생성
- 처리 과정:
- 이미지 전처리 (MiDaS transform 적용)
- 모델을 통한 깊이 예측
- 예측 결과를 원본 이미지 크기로 보간
- 깊이 맵 정규화
3. 세그멘테이션 (segment
)
- 주요 처리 단계:
- 이미지 로드 및 전처리
- 고해상도 처리를 위한 스케일링 (최소 1024px)
- 깊이 추정 수행
- SAM2를 통한 마스크 생성
- 마스크 후처리 및 정제
- 결과를 원본 크기로 복원
4. 마스크 정제 (refine_mask_with_depth
)
- 모폴로지 연산을 사용하여 마스크 개선
- closing 연산으로 작은 구멍을 제거하고 연결성 강화
5. 시각화 (visualize_results
)
- matplotlib을 사용하여 결과 시각화
- 세 가지 이미지 출력:
- 원본 RGB 이미지
- 추정된 깊이 맵
- 세그멘테이션 결과 (컬러 오버레이)
메인 실행 로직
- 장치 선택 (CUDA/CPU)
- SegmentAnyRGBD 인스턴스 생성
- 지정된 디렉토리의 이미지들을 순차적으로 처리
- 깊이 맵을 컬러 히트맵으로 변환하여 저장
사용된 주요 라이브러리
- OpenCV (cv2)
- PyTorch
- NumPy
- Matplotlib (시각화용)
반응형
'Development > SAM' 카테고리의 다른 글
[SAM to OBB] Generate Point coord Tool (0) | 2024.11.11 |
---|---|
[SAM to OBB] Generate Point coord Tool (0) | 2024.11.10 |
Oriented Bounding Boxes Label Editor (0) | 2024.11.08 |
SAM2 (Segment Anything Model 2)에서 포인트를 이용한 세그멘테이션 (0) | 2024.11.04 |
SAM (OBB) 객체 추출 (0) | 2024.10.19 |