Improving Pixel Embedding Learning through Intermediate Distance Regression Supervision for Instance Segmentation

이 논문은 거리 회귀 모듈을 통해 생성된 시드와 특징을 결합하여 픽셀 임베딩 학습을 개선함으로써 인스턴스 분할의 정확도를 크게 향상시키는 새로운 아키텍처를 제안합니다.

핵심

이 논문은 **"복잡하게 엉켜 있는 나뭇잎이나 세포들을 컴퓨터가 하나하나 정확히 구분해 내는 방법"**을 개선한 연구입니다.

컴퓨터 비전 (이미지 인식) 분야에서 '인스턴스 분할 (Instance Segmentation)'이라는 작업은 마치 한 그릇에 담긴 다양한 색상의 스프aghetti(면) 를 하나씩 분리해 내는 것과 비슷합니다. 면들이 서로 얽혀 있고 모양도 불규칙할 때, 컴퓨터가 "이 면은 A, 저 면은 B"라고 정확히 구분하는 것은 매우 어렵습니다.

이 연구는 그 난이도를 낮추기 위해 **"중간 단계의 지도 (Intermediate Supervision)"**라는 새로운 전략을 도입했습니다. 아래에 쉬운 비유로 설명해 드리겠습니다.

---

1. 기존 방식의 문제점: "혼란스러운 방"

기존의 컴퓨터 모델 (U-Net 등) 은 이미지 속의 모든 픽셀 (화소) 을 고차원 공간으로 옮겨서 "이 픽셀은 같은 물체야, 저 픽셀은 다른 물체야"라고 배웁니다.

• 비유: 마치 어두운 방에 수많은 사람 (픽셀) 을 넣고, 같은 팀끼리 모여 있게 하라고 지시하는 것과 같습니다.
• 문제: 사람들이 서로 너무 가깝게 붙어 있거나 (나뭇잎이 겹친 경우), 팀 구분이 모호하면 컴퓨터가 "누가 누구인지" 혼란을 겪어 실수를 많이 합니다.

2. 이 연구의 핵심 아이디어: "지도 없는 나침반"

저자들은 **"가장 어려운 문제 (팀 나누기) 를 바로 해결하려 하지 말고, 먼저 쉬운 문제를 풀게 하자"**는 생각에서 출발했습니다.

• 쉬운 문제: "이 픽셀이 물체의 **가장자리 (경계)**에서 얼마나 떨어져 있는가?"를 계산하는 것입니다.
• 비유: 물체 (나뭇잎) 의 가장자리는 검은색 테두리이고, 중심부는 흰색이라고 칩시다. 컴퓨터에게 "테두리는 검고, 중심은 밝아"라고 가르치는 것은 매우 쉽습니다. 마치 등대가 바다의 가장자리를 비추는 것과 같습니다.

이 연구는 이 "등대 (가장자리 거리 정보)"를 먼저 학습시킨 후, 그 정보를 이용해 복잡한 "팀 나누기"를 하도록 했습니다.

3. 새로운 방법론: "W-Net" (두 단계 학습)

저자들은 기존의 단일 모델 대신 **두 단계를 거치는 'W-Net'**이라는 구조를 만들었습니다.

1. 1 단계 (등대 켜기): 먼저 이미지에서 "물체의 경계까지 얼마나 떨어져 있는가?"를 계산하는 거리 회귀 (Distance Regression) 모듈을 작동시킵니다. 이때 컴퓨터는 "여기가 경계야, 저기는 중심이야"라는 **간단한 지도 (Distmap)**를 그립니다.
2. 2 단계 (팀 나누기): 이제 이 **지도 (거리 정보)**를 원래 이미지와 합쳐서 (붙여서) 두 번째 모듈에 줍니다.
   • 비유: 이제 컴퓨터는 "어두운 방"에 들어가는 것이 아니라, **"등불이 켜진 방"**에 들어가는 것입니다. "아, 이 픽셀은 경계 근처에 있구나, 저 픽셀은 중심이네"라는 힌트를 받으니, 팀을 나누는 일이 훨씬 수월해집니다.

4. 왜 이렇게 하면 더 잘될까? (중요한 통찰)

• 쉬운 것부터 배우기 (Curriculum Learning): 인간이 복잡한 수학 문제를 풀 때, 먼저 기초 공식을 익히는 것과 같습니다. 컴퓨터도 "경계 찾기"라는 쉬운 과제를 먼저 해결하면, 그 과정에서 배운 특징 (Feature) 이 "물체 구분하기"라는 어려운 과제를 풀 때 큰 도움이 됩니다.
• 결과: 이 방법을 적용하자, 기존 모델보다 성능이 8% 이상 크게 향상되었습니다. 특히 나뭇잎이 빽빽하게 모여 있는 이미지에서도 각 잎을 정확히 분리해 냈습니다.

5. 결론: "나침반을 먼저 챙겨라"

이 논문은 **"복잡한 문제를 풀기 전에, 그 문제의 핵심적인 단서 (경계 정보) 를 먼저 찾아주는 것이 중요하다"**는 것을 증명했습니다.

• 기존: "이게 뭐야? (직접 분류)"
• 이 연구: "이게 어디에 있나? (거리 파악) → 아하, 그럼 이게 뭐야? (분류)"

이처럼 **중간 단계의 힌트 (거리 회귀)**를 활용함으로써, 컴퓨터가 나뭇잎이나 세포처럼 복잡하고 얽힌 사물들을 훨씬 더 똑똑하게 구분해 낼 수 있게 되었습니다. 이 기술은 농업 (작물 분석) 이나 의학 (세포 counting) 분야에서 매우 유용하게 쓰일 것으로 기대됩니다.

기술 요약

1. 문제 정의 (Problem)

• 인스턴스 분할의 어려움: 식물 형태 분석 (Plant Phenotyping) 및 세포 정량화와 같은 생물학적 응용 분야에서 개별 객체 (잎, 세포 등) 를 정확히 분할하는 것은 중요합니다. 특히 밀집된 객체나 복잡한 형태를 가진 객체의 경우 기존 경계 상자 (Bounding Box) 기반 방법 (예: Mask R-CNN) 보다 픽셀 임베딩 (Pixel Embedding) 학습과 클러스터링을 통한 접근 방식이 유리합니다.
• 기존 방법의 한계: 픽셀 임베딩 학습은 동일한 객체의 픽셀은 임베딩 공간에서 가깝게, 다른 객체는 멀리 있도록 학습합니다. 그러나 기존 연구들 (CVPPP Leaf Segmentation Challenge 등) 에서 학습된 임베딩 공간이 최적화되지 않아, 특히 잎의 경계와 중맥 (midvein) 이 모호한 부분이나 밀집된 객체 간의 분할에서 성능 저하가 발생했습니다.
• 학습 난이도: 임베딩 학습 자체는 상대적으로 복잡하고 어려운 작업인 반면, 객체 경계까지의 거리를 예측하는 '거리 회귀 (Distance Regression)' 작업은 상대적으로 학습이 쉽습니다.

2. 제안 방법론 (Methodology)

저자들은 W-Net이라고 명명한 새로운 2 단계 아키텍처를 제안하며, 이는 기존 2 개의 헤드를 가진 U-Net 을 개선한 것입니다.

핵심 구조 및 프로세스

1. 거리 회귀 모듈 (Distance Regression Module):

   • 첫 번째 U-Net 을 사용하여 입력 이미지의 픽셀에서 객체 경계까지의 거리 맵 (Distmap) 을 학습합니다.
   • 이 모듈은 MSE 손실 함수로 학습되며, 객체의 윤곽을 파악하는 '시드 (Seed)' 후보를 생성하는 역할을 합니다.
   • 학습된 거리 특징 (Distance Features, D-feat.) 은 객체 존재 확률 (Object-ness) 과 위치 정보를 내포하고 있습니다.

2. 중간 거리 회귀 감독 (Intermediate Distance Regression Supervision):

   • 제안된 방법의 핵심은 **거리 회귀 모듈에서 학습된 특징 맵 (D-feat.) 을 두 번째 U-Net (임베딩 모듈) 의 입력으로 연결 (Concatenation)**하는 것입니다.
   • 단순히 거리 맵 (Distmap) 을 사용하는 것뿐만 아니라, 32 차원의 거리 특징 맵을 입력 이미지에 병합하여 임베딩 학습을 돕습니다. 이는 임베딩 학습이 어려운 경계 영역이나 밀집 영역에서 추가적인 위치 및 형태 정보를 제공합니다.

3. 임베딩 모듈 (Embedding Module):

   • 두 번째 U-Net 은 입력 이미지와 거리 특징을 결합하여 픽셀 임베딩을 생성합니다.
   • 손실 함수: 코사인 유사도 (Cosine Similarity) 기반의 손실 함수를 사용하며, **지역 제약 (Local Constraints)**을 도입합니다.
     • $L_{inter}$: 서로 다른 객체의 임베딩을 분리 (Orthogonal).
     • $L_{intra}$: 동일한 객체 내 픽셀 임베딩을 응집.
     • 지역 제약의 이점: 모든 객체를 전역적으로 구별할 필요 없이, 인접한 객체들만 구별하면 되므로 저차원 임베딩 공간에서도 효율적인 학습이 가능합니다.

4. 클러스터링:

   • 거리 맵의 로컬 최대값 (Local Maxima) 을 시드로 추출한 후, 임베딩 벡터 간의 각도 (Angular Clustering) 를 기반으로 빠른 클러스터링을 수행하여 최종 인스턴스 분할 결과를 도출합니다.

3. 주요 기여 (Key Contributions)

1. W-Net 아키텍처 제안: 거리 회귀 특징을 중간 단계에서 임베딩 학습 모듈로 전달하여 픽셀 임베딩 학습의 정확도를 획기적으로 개선하는 새로운 네트워크 구조를 제안했습니다.
2. 성능 향상: CVPPP Leaf Segmentation Challenge 에서 기존 설정 (거리 특징 연결 없이) 대비 mSBD 점수가 8% 이상 (0.794 → 0.879) 향상되었으며, CodaLab 리더보드에서 1 위를 차지했습니다.
   • 특히 Arabidopsis 이미지 (A1, A2, A4) 에서는 2 위 팀보다 평균 3% 이상 높은 점수 (0.883 → 0.917) 를 기록했습니다.
3. 실험적 검증 (Ablation Studies):
   • 연결 계층 (Concatenative Layer): 거리 특징 맵 (32 차원) 을 연결하는 것이 단순 거리 맵이나 임베딩 특징을 연결하는 것보다 성능이 우수함을 입증했습니다.
   • 지역 vs 전역 제약: 지역 제약 (Local Constraints) 이 저차원 (8 차원) 임베딩 공간에서도 객체 분리에 효과적임을 보였으며, 전역 제약보다 4% 이상 높은 성능을 보였습니다.
   • 임베딩 차원: 8 차원 임베딩이 최적의 성능을 보였으며, 차원이 증가할수록 학습이 어려워져 성능이 저하되는 경향을 발견했습니다.
   • 손실 가중치: 객체 간 손실 ($L_{inter}$) 과 객체 내 손실 ($L_{intra}$) 의 가중치 비율 ($\lambda=1$) 이 최적임을 확인했습니다.

4. 실험 결과 (Results)

• CVPPP Leaf Segmentation Challenge:
  • 전체 mSBD 점수: 0.879 (리더보드 1 위).
  • Arabidopsis 데이터셋 평균 mSBD: 0.917 (2 위 팀 대비 3% 이상 우위).
  • Tobacco 데이터셋 (A3) 은 학습 데이터의 불균형 (783 장 vs 27 장) 으로 인해 상대적으로 낮은 점수 (0.77) 를 기록했으나, 이는 데이터 양의 중요성을 시사합니다.
• Human U2OS Cells (BBBC006v1):
  • 세포 분할 테스트에서도 W-Net 이 U-Net 대비 mSBD (0.896 → 0.915) 와 mAP (0.577 → 0.664) 모두에서 유의미한 개선을 보였습니다. 특히 경계 부위의 불완전한 분할 문제가 해결되었습니다.
• 클러스터링 비교: Angular Clustering 이 Mutex Watershed, Mean Shift, HDBSCAN 보다 성능과 속도 면에서 우수함을 확인했습니다.

5. 의의 및 결론 (Significance & Conclusion)

• 쉬운 작업에서 어려운 작업으로의 전이 학습: 상대적으로 학습이 쉬운 '거리 회귀' 작업에서 얻은 특징을 '인스턴스 분할 (임베딩 학습)'이라는 어려운 작업의 중간 단계에 주입함으로써, 전체 학습 과정을 효과적으로 촉진했습니다.
• 효율적인 임베딩 공간 활용: 지역 제약 (Local Constraints) 과 코사인 임베딩 손실의 결합을 통해 저차원 임베딩 공간에서도 복잡한 객체 분리가 가능함을 입증했습니다.
• 간단하지만 강력한 접근: 복잡한 아키텍처 변경 없이, 기존 U-Net 기반 구조에 중간 감독 (Intermediate Supervision) 을 추가하는 간단한 방식으로 SOTA(State-of-the-Art) 성능을 달성했습니다.

이 논문은 픽셀 임베딩 기반 인스턴스 분할 분야에서 중간 단계의 보조 작업 (거리 회귀) 을 어떻게 활용하여 주 작업의 성능을 극대화할 수 있는지에 대한 중요한 통찰을 제공합니다.