Kaggle Biohub 챌린지 회고: 세포 추적 및 계보 그래프 생성

이번에 Kaggle Biohub 대회에 참가했습니다. 개인 순위는 아쉽게 메달권에는 들지 못했습니다.

따라서 다음 기회에 활용하기 위해 상위권 팀들의 write-up 방식을 정리했습니다.

이번 요약에서는 각 write-up의 세부 사항을 정리하는 것이 아니라, 사용된 접근 방식(approach), 기술(technology), 방법론(methodology)을 정리하고 이해하는 것을 목적으로 합니다.

• Biohub - Cell Tracking During Development

• 과제: 제브라피쉬 배아의 3D 타임랩스 현미경 영상에서 세포핵을 검출하고, 프레임 간에 연결하여 분열까지 포함한 계보 그래프(lineage graph)를 출력합니다.

• 입력: 영상 1개, 100 프레임. 각 프레임은 64 × 256 × 256 voxel (z 방향만 해상도가 낮음).

• 출력: 세포 1개당 1 노드와, 프레임 간의 에지(edge). 자식을 2개 가진 노드는 분열로 처리됩니다.

• 평가 지표: 노드 수로 보정된 에지 Jaccard + 0.1 × 분열 Jaccard.

• 주석이 없는 세포 간의 에지는 기본적으로 감점되지 않습니다. 다만, 추정 세포 수보다 많은 노드를 출력할 경우 감점됩니다.

• 데이터 특징

주석(annotation)은 약 13만 3천 개 노드로, 전체 세포의 1~3%에 불과하며, 분열은 학습 데이터 전체에서 단 151건입니다. - 학습용 영상은 두 개의 배아에서 잘라낸 것입니다. public LB와 private LB는 각각 학습에 사용되지 않은 다른 배아이며, 세포 밀도 역시 크게 다릅니다.

• 학습 데이터에는 같은 프레임의 중복이나, 이미지 전체가 크게 벗어나는 부분(drift)이 있습니다.

이 대회에서는 tracking-by-detection (먼저 세포를 검출하고, 검출점들끼리 연결하는 방식)이 표준적인 형태로 자리 잡았으며, 상위 5개 팀의 해결책도 대략적인 흐름은 공통적입니다.

• 검출(Detection): 각 프레임에서 세포핵의 중심을 찾습니다.
• 연결(Link): 인접한 프레임의 검출점들 간 대응 관계를 점수화합니다.
• 분열 판정(Division Determination): 어떤 세포가 분열할지, 딸세포는 무엇인지 판정합니다.
• 전체 결정(Overall Decision): ILP / LP (정수 계획/선형 계획)로 종합 최적화하거나, greedy (좋은 후보부터 순차적으로 확정) 방식으로 트랙 전체를 결정합니다.
• 후처리(Post-processing): 갭 보완, 짧은 트랙 삭제, 좌표 평활화 등을 수행합니다.

• 1st Place Solution
• 한마디로 말하자면: 직접 만든 특징에만 의존하지 않고, 3프레임의 국소 이미지로부터 대상 세포의 분열 상태와, 전후 프레임에서의 동일 세포 및 모녀 위치를 추정하는 모델(Soon Net)을 만들고, 그 출력을 기하학적 특징과 결합하여 학습 링크어에 전달했습니다.

파이프라인: 검출 (3D U-Net) → Soon Net → 학습 링크어 → greedy decode → 후처리

검출: 3D U-Net + Cellpose 스타일 헤드

• 전경 확률과 '세포 중심으로의 흐름(flow)'을 출력하고, 이 흐름을 따라 같은 중심으로 모인 voxel들을 하나의 세포로 그룹화합니다 (Cellpose 방식).

• 마스크가 얻어지므로, 부피나 모양을 링크어 특징에 활용할 수 있습니다.

• 학습 과정:

• 사전 학습: 가벼운 MAE (입력의 일부를 숨기고 복원시키는 자기 지도 학습). 세포 전체가 가려지지 않도록 작은 블록으로 마스킹합니다.

• 학습: 직접 그린 마스크와 GT 노드에 찍은 작은 스탬프를 교사(teacher)로 사용합니다.

• 재학습: 2의 예측을 검출 신뢰도와 트랙 지속 길이로 제한하여 가상 레이블(pseudo label)로 사용합니다.

• 외부 데이터: Ultrack (세포 추적 도구)의 공개 제브라피쉬 배아 데이터도 사전 학습과 가상 레이블에 의한 재학습에 사용했습니다.

Soon Net: 세포별 분열 상태 및 대응 위치 예측 (이 해결책의 핵심)

• 입력: 세포별로 t−1, t, t+1의 3프레임을 잘라내고, t의 대상 세포 위치를 나타내는 Gaussian 표시를 추가합니다. 이를 통해 '어떤 세포에 대해 답할지'를 지정합니다. t+1

이미지도 보고, 미지의 미래를 예측하는 것이 아니라 이전/이후 이미지 내에서 대응되는 세포를 찾는 방식입니다. - 구성: 작은 3D CNN으로 토큰화한 후 Transformer에 통과시킵니다. Primus의 설계를 참고하여 CNN tokenizer를 소형화하고, Transformer에는 EVA-02 블록을 사용합니다.

• 출력: 분열 상태의 3개 클래스(간기/곧 분열할 것/분열 직후)와 FPN 헤드에서 각 프레임에 대해 내보내는 2종류의 occupancy map입니다.

• Continuation map: 동일 세포의 전후 프레임에서의 위치를 나타냅니다.

• Division map: 모세포에서 본 딸 2개의 위치나, 딸에서 본 모세포의 위치를 나타냅니다.

• 학습 시工夫:

• 잘라내는 위치를 크게 벗어나게 합니다. 중앙 세포만 답하는 치팅을 막고 Gaussian 표시자를 사용하게 합니다.

• 간기 세포에는 시간 방향의 augmentation을 사용합니다.
  t−1
  대신 t−2, t−3,
  t+1
  대신 t+2, t+3도 사용하여 큰 이동을 학습시킵니다. - 초기 모델로 분열 후보를 찾아내고, 육안으로 확인하여 학습 데이터를 늘립니다.

• OOF 예측에서 간기인데 분열 전후로 오판된 hard negative를 수집합니다.

링커(Linker): LightGlue 스타일의 Transformer

• LightGlue(2장의 이미지 특징점을 대응시키는 Transformer)의 아이디어를 참고합니다. 부피, 형태, 위치 등의 노드 특징으로 각 검출을 토큰화하고, 프레임 t와 t+1의 근접 세포들 사이에 cross-attention을 겁니다. - 쌍 특징에는 이동량이나 형태 변화 외에도 Soon Net의 occupancy map을 상대방 위치에서 읽은 값을 사용합니다. 부모의 다음 프레임 맵을 자녀 후보 위치에서, 자녀 후보의 이전 프레임 맵을 부모 위치에서 읽어, 양방향 대응의 강도를 살펴봅니다.
• Pair Head: 두 세포의 임베딩과 쌍 특징으로부터 각 후보 링크를 점수화합니다.
• Fork Head: 각 부모에 대해 '자식 없음/자식 1개/자식 2개'를 평가합니다. 분열은 '분열 여부(kind)'와 '어느 것이 자식인지(who)'로 나누고,
  p(kind) × p(who | kind)
  으로 계산합니다. 하나의 softmax에 모으면 다수의 자식 조합 속에 희귀한 분열 판단이 묻혀 오분열이 증가했기 때문입니다.

디코드 및 후처리

• ILP는 사용하지 않고 greedy decode를 합니다. 각 부모의 선택지(자식 1개에 연결, 분열)를 '자식 없음'과 비교하여, '자식 없음'보다 더 유력한 것만 차이가 큰 순서대로 채택합니다.

• 각 부모는 최대 1개의 선택지를, 각 자식은 최대 1개의 부모라는 제약을 지킵니다. 부모가 할당되지 않은 세포는 새로운 트랙의 시작점으로 삼습니다.

• 후처리는 최소화합니다: 분열을 넘어서 딸이 모 쪽으로 끌려가는 line-fit 평활화, 태어난 자식의 다음 프레임에서의 재분열 금지, 1프레임 간격 보완, 3프레임 미만의 트랙 삭제.

• 2위 솔루션

• 한마디로 말하면:
  약한 검출도 바로 버리지 않고, 시간 방향으로 연결되는지 여부로 판단한다.

파이프라인: temporal 3D U-Net (주 모델 + 보조 모델) → 후보 선택/링크/분열 판정 → 보조 트랙으로 복구 → 희미한 트랙의 복원

검출: temporal 3D U-Net

• 전처리: XY만 블록 평균으로 1/4로 축소하여 Z와 같은 간격(등방)의 볼륨으로 만듭니다.

• 모델:
  t−1, t, t+1
  의 3프레임을 넣는 residual 3D U-Net (5프레임 버전도 병용). - 출력은 4가지입니다.

• 중심 히트맵: 세포 중심 후보지.

• 오프셋: 축소된 거친 그리드 내에서 중심 위치를 세밀하게 보정합니다.

• 이전 프레임으로의 움직임(불확실성 포함): 이전 프레임의 어떤 세포로부터 왔는지 후보를 내는 데 사용합니다.

• 특징 벡터(descriptor): 세포의 외형을 나타내는 벡터. 링크 판정에 사용합니다.

• 희소 라벨 대책: '어노테이션 있음/이미지에서 배경으로 판단 가능한 영역/불명'으로 나누고, 불명한 영역은 히트맵의 loss에서 제외합니다.

• 합성된 희미한 세포: 어노테이션이 있는 실제 세포의 짧은 트랙을 잘라내고 배경을 빼낸 후 어둡게/흐리게 하여 같은 영상의 다른 장소에 붙입니다. - TTA 등으로 평균하면 사라져 버리는 피크도, 원래 방향 예측으로부터 후보로 되돌립니다.

• 보조 검출기는 주 검출기와 단순히 평균하지 않고, 별도의 트랙을 만들어서 빠진 트랙의 보완과 링크/분열 확인에 사용합니다.

트래킹

• 후보 선택: 검출기가 남기고자 하는 개수보다 많이 후보를 내놓으므로, 그중에서 남길 세포를 고른다.

• 세포 수 추정 모델: 프레임마다 남아있을 세포 수의 대략적인 것을 예측한다.

• LightGBM 2개: 후보 각각에 대해 '어두운 세포인지', '세포다운지'를 추정하여 순위를 조정하고, 기존 후보군을 남겨둔 채로 기존 후보와 떨어진 위치에 있는 후보를 추가한다.

• 최종적인 선택은 검출의 확신도뿐만 아니라, 전후 프레임의 세포와 연결될 가능성도 보고 결정한다 (가까운 중복 제외). 따라서 어둡더라도 전후와 연결되는 점이 밝지만 고립된 점보다 유리할 때가 있다.

• 링크: 검출기가 예측한 '전 프레임으로의 움직임'을 바탕으로, 이전 프레임의 어떤 세포에서 왔는지 후보를 내놓는다. 이를 학습한 스코어러가 거리・움직임・검출의 확신도・트랙의 전후 관계 등을 종합하여 점수를 매긴다.

• 분열: 부모 1개와 딸 후보 2개의 조합을, 위치 관계・외형・전후 프레임의 상황 등에서 종합적으로 점수화하는 별도의 스코어러이다. 기존 트랙과의 경쟁도 고려한다.

• 복구(修復): 보조 검출기에서 만든 트랙을 사용하여 주 트랙에 빠진 경로를 채우고, 링크와 분열이 양쪽 모두 일치하는지 확인한다. 1프레임의 누락까지 연결한다.

촬영 조건별 그룹 분류

• 밝기・대비 등 이미지 통계량을 표준화하여 비디오를 K-means로 2개 그룹으로 나눈다 (어노테이션이나 배아 정보는 사용하지 않는다).
• 한쪽은 배경이 밝고, 약한 세포의 피크가 묻히기 쉽다. 후속 보정 중 일부는 검증에서 효과를 본 그룹에만 사용한다.

핵심 포인트: 희미한 트랙 복원

• 트래커가 버린 후보 중에서 길게 이어지는 경로를 찾아, 원본 이미지에서 '후보 중심이 주변 배경보다 밝은지'를 확인하여 복원한다. 검출기는 다시 돌리지 않는다.

• 배경이 밝은 쪽 그룹에만 적용한다.

• 3위권 솔루션

• 한마디로 말하자면:
  크기가 큰 모델들의 앙상블로 검출을 강화하고, 흐름(flow)・대응(matching)・분열 각 모델의 출력을 LP/MILP를 이용해 종합적으로 최적화한다.

파이프라인: 검출 → dense flow → matcher → 분열 판정 → LP/MILP → 후처리

검출: 2.5D U-Net + 3D SegResNet

• 2.5D U-Net: 인접한 3장의 Z 슬라이스를 채널로 쌓고, 사전 학습된 2D 인코더(EfficientNetV2-L, EfficientNet-B7)로 특징을 추출하여 Z 방향으로 종합해 3D 디코더에 전달한다. 이 3장은 시간 방향이 아니라 깊이 방향의 정보이다. CZII 콘페런스 4위권에서 사용했던 자신들의 모델 구조를 재활용했다.
• 3D 쪽은 MONAI의 SegResNet을 사용한다. 3종류 모델의 히트맵을 가중치 평균하여 사용한다.
• 희소 레이블 대책: DoG(고전적인 blob 검출) 후보 중, 어노테이션과 일치하지 않는 주변 영역을 loss에서 제외시킨다. 정답 예시로 추가하는 것이 아니라, 세포일지도 모르는 장소를 배경으로 학습시키지 않기 위한 처리이다.

Dense flow: 프레임 간 3D 변위장

• 연속된 2프레임 이미지로부터 각 위치가 다음 프레임에서 얼마나 움직일지를 추정하는 작은 3D encoder-decoder이다. 실제 데이터에서는 위치 맞춤 후 이미지의 일관성 등을 이용한 자기 지도 학습(self-supervised learning)을, 변형이 알려진 합성 데이터에서는 지도 학습을 수행한다.
• matcher의 후보 탐색과 분열 모델에 넣을 이미지의 위치 맞춤에 사용된다.

Matcher: 대응(matching) Transformer

• 인접한 2프레임 검출점에 대해 '어느 것과 어느 것이 같은 세포인지'를 점수화한다. 여기서는 후보들을 채점하고, 최종적인 연결은 후속 단계에서 결정한다.
• matcher 전용의 2.5D U-Net 특징을 각 검출점 위치에서 추출하여, flow로 보정된 위치 정보와 함께 프레임 내(self) 및 프레임 간(cross) 어텐션을 이용해 비교한다.
• '대응하는 세포가 없음'(null)에 대한 점수도 낸다.
• 학습 시에는, 어노테이션이 없는 세포를 '대응 상대가 없음(null)'의 정답으로 처리하지 않는다. 어노테이션이 없다고 해서 실제로는 대응하는 세포가 없을 수도 있기 때문이다.

LP/MILP: 링크와 분열을 동시에 선택

• 먼저 링크를 확정하면, 올바른 딸세포가 다른 부모에게 할당될 수 있다. 따라서, 부모 중복 등을 허용하지 않는 제약 조건 하에서, 남길 노드, 일반적인 링크, 그리고 두 개의 딸세포로 구성된 분열 이벤트를 동시에 선택한다.
• 희소 라벨(sparse label) 대책으로, 후보마다 '평가 대상이 될 확률 $a$'와 '평가 대상이었을 경우의 정확도 $q$'를 추정한다. 기대 TP는 $a imes q$, 기대 FP는 $a imes (1 - q)$로 설정하고, 이 값들과 노드 수로부터 공식 스코어의 근사 목적 함수를 만든다. - 솔버는 HiGHS(LP 완화(relaxation)를 풀고, 필요하면 MILP)를 사용한다.

후처리

• 갭 보완(Gap filling), 작은 연결 성분 제거, 아핀 움직임 모델을 사용한 좌표 보정. 좌표 보정 시에는 링크를 변경하지 않고 위치의 흔들림을 줄인다.

• 4th Place Solution

• 한마디로 말하자면: 계산 자원이 부족하므로, **생물학적 관찰(분열의 모양, 핵의 크기, 조직의 이동)**을 작은 모델에 녹여 넣는다. 증거를 조합하는 부분은 CPU 기반의 그래디언트 부스팅이다.

파이프라인: 검출(3D Net) → 링크 확률(LightGBM) → 분열 점수 → ILP → 후처리

검출: Cellpose 스타일 벡터장 3D Net

• 세포 확률과 핵 중심 방향의 단위 벡터를 예측하고, 이 벡터를 따라가며 핵을 찾는다 (1st와 동일하게 Cellpose의 개념).
• dual-encoder 기반의 3D U-Net. 한쪽에는 현재 프레임을, 다른 쪽에는 '다음 프레임 - 이전 프레임'의 차분을 입력한다.
• 세포 확률 손실(loss)은 '어노테이션 있음/확실한 배경/그 외' 세 단계로 가중치를 다르게 적용한 BCE를 사용한다. '그 외'도 완전히 무시하지 않고, 아주 약하게 배경 쪽으로 기울인다. 벡터장에는 L1 loss를 사용한다.
• 유사 라벨(pseudo label)로 학습된 모델(ZebraHub의 어노테이션 없는 배아도 사용)은, 절반 이상의 노드가 GT(Ground Truth)만으로 학습한 모델에서도 뒷받침되는 트랙만을 채택한다.
• DoG를 이용해 영상별 핵 간격을 측정하고, 희소한 영상에서는 검출 임계치를 엄격하게 한다.

링크: LightGBM의 엣지 모델

• 특징(feature)은 기하학적 특징, 공개 베이스라인 구조로 재학습된 링크어의 확률, InfoNCE(대조 학습의 loss)로 학습한 세포 임베딩의 유사도이다.
• 이 해법에서는 OOF(Out-of-Fold) 검증에서 거의 가치가 없어 보였던 Transformer의 특징이, 전체 데이터로 학습시킨 버전에서는 LB 개선에 효과를 보였다. 폴드별 성능만으로는 전체 데이터 학습 버전 특징의 가치를 과소평가하는 경우가 있었다.

분열: 노드별 분열 비용을 ILP에 포함 (최대 개선)

• 분열 비용이 전체적으로 하나일 경우, '분열을 전혀 열지 않거나' 또는 '두 세포가 가까우면 어디서든 연다'는 문제가 발생한다. 따라서 노드마다 비용 $6 - 16 imes s$ (희소 영상에서는 $5 - 16 imes s$)를 설정한다. 여기서 $s$는 해당 노드가 다음에 분열할 모세포일 확률이다. - $s$는 분열 CNN(5프레임 슬라이스 관찰)과 CatBoost/LightGBM의 조합을 사용한다. 특징은 '모세포가 둥글고 밝아지면서, 두 딸세포가 나타나 멀어지는' 생물학에 맞춰 설계한다. - 라벨은 어노테이션된 트랙에서만 가져온다. hard negative mining은 악화되었다 (대부분이 어노테이션 없는 실제 분열이었기 때문).
• ILP 이후 fork verifier로 두 딸세포의 움직임을 확인하고, 분열이 의심스러우면 약한 쪽 딸세포의 엣지를 제거한다.

ILP 및 후처리

• tracksdata(트래킹용 라이브러리) + SCIP(MILP 솔버).

• 움직임 보정된 평활화, 갭 보완, 희소한 영상에서는 다른 검출기에 의한 뒷받침이 적은 트랙 구간 제거.

• 5th Place: 3D U-Net + Transformer Linker + 다단계 ILP Tracking

• 한마디로 말하자면: 공개 베이스라인의 temporal 3D U-Net에 Transformer Linker를 결합하고, ZebraHub에서 사전 학습 및 신규 세포/분열 헤드, 다단계 ILP로 강화했다.

파이프라인: 3D U-Net → Transformer Linker → 다단계 ILP → 후처리

모델

• 검출: 공개 베이스라인의 temporal 3D U-Net (전후 프레임도 입력하는 3D U-Net). XY를 평균 풀링으로 1/4로 축소하여 입력한다.
• 링크어: 프레임 $t$와 $t+1$}**

세포 토큰에 self/cross attention을 적용합니다. 총 4가지를 학습하며, 출력은 링크 스코어, 신규 세포, 분열의 세 가지입니다. 신규 세포는 '이전 프레임에서 대응하는 상대가 없는 세포'를 의미합니다. 나머지 동일성은 U-Net의 특징으로부터 생성된 16차원 외형 벡터로 학습하여 토큰 입력에 넣어 근처의 다른 세포를 구별합니다. - 사전 학습: ZebraHub 공개 데이터로 사전 학습한 후, 대회 데이터로 fine-tune 합니다.

다단계 ILP (Iterative Label Propagation)

검출 신뢰도, 링크 스코어, 동일성 벡터 유사도를 비용에 반영합니다. 신규 세포 점수가 높으면 트랙 시작을, 분열 점수가 높으면 분기 선택을 용이하게 합니다.

• ILP1: 신뢰도가 높은 세포만으로 해결합니다.
• ILP2: 세포를 늘려서 다시 해결합니다. 간극은 채워지지만, 가짜 분열도 증가합니다.
• 새로운 분열 취소: 1에 없었던 분열을 취소합니다. 추가된 세포 자체는 삭제하지 않고, 다른 트랙으로 남겨둡니다.
• ILP3: 세포와 분열은 고정하고, 링크만 다시 해결합니다.
• 복구 및 평활화: ILP4로 떨어진 트랙의 끝점과 시작점을 연결하여 빈 곳을 보완하고, 모호한 링크는 외형으로 재매칭합니다. 마지막으로 트랙 위치를 평활화합니다.