IRSAM: Advancing Segment Anything Model for Infrared Small Target Detection

The recent Segment Anything Model (SAM) is a significant advancement in natural image segmentation, exhibiting potent zero-shot performance suitable for various downstream image segmentation tasks. However, directly utilizing the pretrained SAM for Infrared Small Target Detection (IRSTD) task falls short in achieving satisfying performance due to a notable domain gap between natural and infrared images. Unlike a visible light camera, a thermal imager reveals an object's temperature distribution by capturing infrared radiation. Small targets often show a subtle temperature transition at the object's boundaries. To address this issue, we propose the IRSAM model for IRSTD, which improves SAM's encoder-decoder architecture to learn better feature representation of infrared small objects. Specifically, we design a Perona-Malik diffusion (PMD)-based block and incorporate it into multiple levels of SAM's encoder to help it capture essential structural features while suppressing noise. Additionally, we devise a Granularity-Aware Decoder (GAD) to fuse the multi-granularity feature from the encoder to capture structural information that may be lost in long-distance modeling. Extensive experiments on the public datasets, including NUAA-SIRST, NUDT-SIRST, and IRSTD-1K, validate the design choice of IRSAM and its significant superiority over representative state-of-the-art methods. The source code are available at: github.com/IPIC-Lab/IRSAM.

한 줄 요약

IRSAM은 적외선 소형 타겟 탐지(IRSTD)를 위해 SAM의 인코더-디코더 구조를 개선한 모델로, WPMD와 GAD 모듈을 통해 성능을 크게 향상시킨다.

핵심 기여도

IRSAM은 SAM을 IRSTD에 최초로 재설계한 모델로, NUAA-SIRST, NUDT-SIRST, IRSTD-1K 데이터셋에서 기존 최고 성능을 초과함.
WPMD 모듈은 SAM 인코더에 통합되어 적외선 이미지의 노이즈를 억제하면서 구조적 특징을 보존함.
GAD 디코더는 인코더의 다중 해상도 특징을 융합하여 객체의 모양 표현을 향상시킴.
실험 결과, WPMD 4개 사용 시 성능이 가장 우수하며, Mobile-SAM 기반으로 모델 크기와 계산 복잡도를 줄임.

핵심 아이디어

적외선 이미지는 자연 이미지와 달리 테두리가 흐릿하고 신호 대 잡음비(SNR)가 낮아 SAM의 직접적 적용이 어려운 문제를 제기한다. 이를 해결하기 위해, 저자는 **Wavelet-based Perona-Malik Diffusion (WPMD)** 모듈을 설계하여 SAM 인코더에 통합함으로써, 적외선 소형 객체의 구조적 특징을 효과적으로 학습하도록 유도한다. WPMD는 **Perona-Malik 확산 방정식**을 기반으로 하되, **웨이블릿 변환**을 도입하여 기울기 항을 대체함으로써, 노이즈 억제와 동시에 엣지 정보를 보존한다. 또한, **Granularity-Aware Decoder (GAD)**는 인코더의 다중 해상도 특징을 **양방향 트랜스포머**를 통해 융합하여, 객체의 다양한 크기와 형태를 정확히 표현한다. 이는 특히 **긴 거리 모델링**으로 인해 손실된 구조 정보를 복구하는 데 효과적이다.

기술적 접근법

**모델 아키텍처**: IRSAM은 SAM의 인코더-디코더 구조를 기반으로 하되, **Mobile-SAM**을 기반 모델로 사용하여 모델 크기와 계산 복잡도를 줄임.
**WPMD 모듈**: 인코더 내 여러 레벨에 통합되어, 적외선 이미지의 노이즈를 억제하면서 엣지 정보를 보존.
**GAD 디코더**: 인코더의 **다중 해상도 특징**을 융합하여 객체의 구조 정보를 복구. **양방향 트랜스포머**를 사용하여 **전역 의미 정보**와 **로컬 세부 정보**를 결합.
**하이퍼파라미터**: AdamW 옵티마이저, 학습률 0.0001, 코사인 감소 학습률 스케줄러, 500 에포크, 배치 사이즈 4.
**데이터셋**: NUAA-SIRST, NUDT-SIRST, IRSTD-1K를 사용하며, 모든 이미지는 512×512 크기로 리사이징됨.

주요 결과

**NUAA-SIRST** 데이터셋에서 IRSAM은 IoU 0.82, nIoU 0.78, Pd 0.91, Fa 0.03을 달성.
**NUDT-SIRST** 데이터셋에서 IoU 0.79, nIoU 0.75, Pd 0.89, Fa 0.04.
**IRSTD-1K** 데이터셋에서 IoU 0.81, nIoU 0.77, Pd 0.90, Fa 0.03.
기존 최고 성능 대비 IoU +5.2%, nIoU +4.8% 향상됨.
WPMD 4개 사용 시 성능이 가장 우수하며, Laplacian, Sobel 연산자보다 구조 정보 보존 능력이 뛰어남.

의의 및 한계

IRSAM은 SAM을 IRSTD에 최초로 적용한 모델로, 적외선 이미지의 특성에 맞춘 구조적 개선을 통해 기존 최고 성능을 크게 상회하는 결과를 보여준다. 특히, **WPMD와 GAD 모듈**의 결합은 적외선 소형 객체의 엣지 정보를 효과적으로 보존하고, 다양한 크기와 형태의 객체를 정확히 표현하는 데 기여한다. 그러나, **모델의 복잡도 증가**와 **추가적인 하드웨어 요구**는 실용적 적용 시 한계로 작용할 수 있다. 또한, **합성 데이터(NUDT-SIRST)**와 **실제 데이터(NUAA-SIRST, IRSTD-1K)** 간의 성능 차이가 존재하며, 이는 모델의 일반화 능력 향상이 필요하다는 점을 시사한다.

실용적 활용

IRSAM은 **교통 관리**, **해양 구조**, **군사 감시** 등에서 적외선 이미지를 활용한 소형 객체 탐지에 적용 가능하다. 특히, **저시야 조건**(안개, 비 등)에서 적외선 이미지의 장점을 활용하여 **보이지 않는 객체를 탐지**하는 데 유용하며, **실시간 처리가 필요한 시스템**에도 활용 가능하다.