[데이터셋] HEST-1k 사용 방법

HEST-1k 다운로드 및 코드에서 사용하는 방법

 

논문 -> https://arxiv.org/abs/2406.16192

github -> https://github.com/mahmoodlab/HEST/tree/main

huggingface -> https://huggingface.co/datasets/MahmoodLab/hest

---

위 huggingface와 github에 기재된 튜토리얼을 바탕으로, google colab에서 cpu를 사용하였다.

 

1. Huggingface 토큰 생성

HEST-1k 데이터셋을 사용하려면 우선 huggingface에서 생성한 토큰을 통해 코드에서 데이터셋에 접근해야 하며, 아래 설명을 순서대로 수행한다.

 

1-1. Huggingface 계정 생성

설명 생략

 

1-2. 프로필 -> Access Tokens에서 토큰 생성

우측 상단 프로필 -> Access Tokens

Access Tokens 화면에서 토큰 생성하기 선택

토큰 생성 화면에서 이미지 설정대로 토큰 생성

이때 이미지와 같이

 

- Token type=Fine-grained

- 토큰 이름 설정: 원하는 대로

- User permissions/Repositories에서 3개 옵션 모두 체크

 

로 설정하는데, 화면에 표시되지 않은 다른 체크박스들은 모두 비워놓아도 무방하다.

 

토큰 설정을 완료한 다음 페이지 하단에서 "Create token"을 클릭하여 토큰을 생성하면

이미지와 같이 토큰 코드를 저장할 수 있는 팝업이 표시되며, 이 토큰을 복사한 뒤 복사한 토큰을 통해 코드에서 huggingface 레포지토리에 접근하게 된다. (토큰 코드는 나중에도 huggingface>Access token 창에서 확인할 수 있지만 매번 huggingface 사이트에 접속하는 것도 귀찮으니 편한 위치에 따로 저장해 두는 것을 추천)

 

2. 데이터셋 다운로드

이후 colab (또는 각자의 개발 환경)에서 진행한다.

 

2-1. Huggingface 로그인

!pip install huggingface-hub

from huggingface_hub import login

login(token="YOUR HUGGINGFACE TOKEN")	# 아까 복사한 토큰

필요한 라이브러리/모듈들을 설치해 주고 huggingface login을 import한 뒤 아까 복사한 토큰을 사용해 로그인한다.

 

HEST 공식 huggingface 튜토리얼에는 huggingface_hub만 설치하는 것으로 설명되어 있는데, 추후 iter_hest 함수를 통해 샘플 내부에 접근하려면 !pip install huggingface_hub만 하고 실행시켰을 때 오류가 발생할 수 있으므로 오류 로그를 보고 각자의 환경에 맞게 필요한 것을 설치하거나 불필요한 것을 삭제하면 될 듯.

 

참고)

!pip uninstall -y rapids-dask-dependency

!apt-get -y install openslide-tools

!pip install openslide-python

!pip install git+https://github.com/mahmoodlab/HEST.git

(제 개발 환경에서는 iter_hest 사용을 위해 위와 같이 추가로 설정해 주어야 했습니다)

 

2-2. 데이터셋 다운로드

 

이후 다운로드 함수를 구현하여 데이터셋을 다운받아야 하는데,

 

(1) 데이터셋 전체를 다운받는 방법

(2) 일부 데이터셋만 다운받는 방법 (샘플 id or 암종/장기명으로 다운로드)

 

이 있다.

 

다만 전체 데이터셋 용량이 약 1.6TB이므로 개인 작업 환경에서 (1)로 진행하기는 어렵지 않을까 싶고,

일부 데이터셋 다운로드 위주로 설명할 예정이다.

 

우선은 아래와 같이 huggingface로부터 데이터셋을 다운받아 hest_data 폴더 아래로 디렉토리 구조에 맞게 저장해 주는 download_hest() 함수를 구현한다.

import os
import zipfile

from huggingface_hub import snapshot_download
from tqdm import tqdm


def download_hest(patterns, local_dir):
    repo_id = 'MahmoodLab/hest'
    snapshot_download(repo_id=repo_id, allow_patterns=patterns, repo_type="dataset", local_dir=local_dir)

    seg_dir = os.path.join(local_dir, 'cellvit_seg')
    if os.path.exists(seg_dir):
        print('Unzipping cell vit segmentation...')
        for filename in tqdm([s for s in os.listdir(seg_dir) if s.endswith('.zip')]):
            path_zip = os.path.join(seg_dir, filename)
                        
            with zipfile.ZipFile(path_zip, 'r') as zip_ref:
                zip_ref.extractall(seg_dir)


local_dir='hest_data' # hest will be dowloaded to this folder

위 코드 그대로 변형 없이 사용해도 무방하며,

데이터셋이 저장될 폴더명을 변경하고 싶다면 local_dir 변수에 다른 값을 설정하면 된다.

 

2-2-1. 암종/장기명으로 다운로드

 

첫 번째는 암종이나 장기를 지정해 해당 암종 or 장기의 샘플만 다운로드하는 방식이다.

import datasets
import pandas as pd

local_dir='hest_data' # hest will be dowloaded to this folder

meta_df = pd.read_csv("hf://datasets/MahmoodLab/hest/HEST_v1_1_0.csv")

# Filter the dataframe by organ, oncotree code...
meta_df = meta_df[meta_df['oncotree_code'] == 'IDC']	# <- oncotree code로 필터링
meta_df = meta_df[meta_df['organ'] == 'Breast']	# <- 장기로 필터링

ids_to_query = meta_df['id'].values

list_patterns = [f"*{id}[_.]**" for id in ids_to_query]
download_hest(list_patterns, local_dir) # see method definition above

Huggingface에서 HEST 전체 데이터셋을 요약하는 csv파일을 받아 와 meta_df에 저장한 다음,

원하는 암종 코드 또는 장기에 해당하는 샘플의 샘플 id를 ids_to_query 리스트에 저장하고 앞서 구현한 download_hest() 함수를 통해 데이터셋을 다운로드하는 코드이다.

# Filter the dataframe by organ, oncotree code...
meta_df = meta_df[meta_df['oncotree_code'] == 'IDC']	# <- oncotree code로 필터링
meta_df = meta_df[meta_df['organ'] == 'Breast']	# <- 장기로 필터링

부분을 적절히 수정하여 원하는 샘플을 추출할 수 있으며,

oncotree code만 사용하거나, 장기만 사용하거나, 또는 둘 다 사용하여 필터링할 수 있다.

 

이때 oncotree code는 암 유형 식별을 위해 부여한 코드인데,

아래 사이트에서 암종과 해당 암종에 대응하는 코드를 확인할 수 있다.

oncotree code 확인 -> https://oncotree.mskcc.org/?version=oncotree_latest_stable&field=NAME

괄호 안에 있는 문자열이 oncotree code이다.

 

예를 들어, Breast Invasive Lobular Carcinoma 샘플을 사용하고 싶다면 

meta_df = meta_df[meta_df['oncotree_code'] == 'ILC']

와 같이 설정할 수 있다.

 

다만 HEST 데이터셋에 어떤 암종과 기관이 포함되어 있는지 미리 알고 있어야 하며,

매번 oncotree 사이트에 방문하여 코드를 확인하는 것도 귀찮기 때문에,

huggingface 튜토리얼에는 없지만 HEST_v1_1_0.csv 파일을 통해 HEST 데이터셋에 어떤 암종/기관의 샘플들이 몇 개 포함되어 있는지 추출하는 코드를 작성했다.

import pandas as pd

# metadata
meta_df = pd.read_csv("hf://datasets/MahmoodLab/hest/HEST_v1_1_0.csv")

# organ/ontotree code별 개수 카운트
counts = meta_df.groupby(['organ', 'oncotree_code']).size().reset_index(name='count')
print(counts)

코드를 실행하면 아래와 같이 출력된다.

즉, 위 이미지에 포함된 장기나 oncotree code만 필터링에 사용할 수 있다.

 

정리하자면,

download_hest() 함수를 구현하고, meta_df에서 HEST-1k 데이터셋의 구성 정보를 추출한 다음, 존재하는 암종과 장기에 대해서 필터링하여 원하는 샘플을 다운로드하는 파이프라인으로 구성된다.

 

2-2-2. 샘플 id로 다운로드

 

두 번째는 샘플 id를 사용해 어떤 샘플을 다운받을지 하드코딩하는 방식인데,

import datasets

ids_to_query = ['TENX96', 'TENX99'] # list of ids to query

list_patterns = [f"*{id}[_.]**" for id in ids_to_query]
download_hest(list_patterns, local_dir) # see method definition above

 

2-2-1처럼 meta_df에서 암종 코드와 장기를 통해 샘플 id를 자동 필터링하는 것이 아니라,

ids_to_query 리스트에 사용할 샘플 id를 직접 전달하게 된다.

 

사용하고 싶은 암종/장기 샘플의 샘플 id를 미리 확인하려면 

https://huggingface.co/datasets/MahmoodLab/hest/resolve/main/HEST_v1_1_0.csv

링크를 방문하고 해당 csv 파일을 확인해

파일의 img_filename 열에서 '.' 앞의 문자열이 샘플 id이므로 이를 parsing해 사용할 수 있다.

 

25.12.01 수정:

바로 옆에 id 열이 있는데 못 봄(ㅜㅜ)

굳이 parsing할 필요 없고, parsing하는 부분 삭제하고, grouped 정의하는 부분에서 'sample_id' 부분을 'id'로 바꿔주면 된다.

 

그런데 위와 같이 샘플 id를 직접 확인하는 과정이 매우 번거로우므로,

huggingface 튜토리얼에는 없지만 meta_df에서 암종/기관별로 해당 카테고리의 샘플 id를 출력하는 코드를 작성했다.

import pandas as pd

# metadata
meta_df = pd.read_csv("hf://datasets/MahmoodLab/hest/HEST_v1_1_0.csv")

# pd.set_option('display.max_colwidth', 100)   # 한 column 길이 제한
# pd.set_option('display.max_columns', None)    # 모든 column 표시 

# '.' 기준으로 샘플 id 추출
meta_df['sample_id'] = meta_df['image_filename'].str.split('.').str[0]

# organ/oncotree_code 그룹별 count + 해당 그룹의 sample_id list
grouped = (
    meta_df
    .groupby(['organ', 'oncotree_code'])
    .agg(
        count=('sample_id', 'count'),
        sample_ids=('sample_id', list)
    )
    .reset_index()
)

print(grouped)

코드를 실행하면 아래와 같이 출력된다.

샘플 id 기준으로 특정 암종의 샘플 몇 개만 다운받고 싶을 때 사용할 수 있다.

 

3. 샘플 출력

이제 다운받은 샘플 내부 데이터에 접근한다.

from hest import iter_hest

for st in iter_hest('../hest_data', id_list=['TENX95']):
    print(st)

이와 같이 HEST 라이브러리의 iter_hest() 함수를 사용해 접근하게 되는데,

colab에서 작업할 경우 iter_hest의 경로 파라미터로 절대경로 '/content/hest_data'를 전달하거나,

상대경로를 사용할 경우 './hest_data'로 전달해야 한다.

 

위 코드를 실행시키면 다음과 같이 출력된다.

 

이때 iter_hest()는 HESTIterator 객체를 생성해 샘플들을 순회하며 HESTData 객체를 로딩하는 함수로,

출력에서 HESTData 객체가 출력된 것을 확인할 수 있다.

 

HESTIterator, HESTData, iter_hest()에 관해서는 HEST github의 /src/hest/HESTSData.py 파일(-> https://github.com/mahmoodlab/HEST/blob/main/src/hest/HESTData.py#L1145) 전체 코드를 확인할 수 있는데,

 

간략히 설명하자면,

iter_hest() 함수를 통해 해당 샘플의

- ST 데이터 (anndata 객체)

- WSI (Whole Slide Image)

- 메타데이터

에 접근할 수 있다.

 

이를 응용하여 HESTData 객체 내부에서 anndata 객체, 이미지, 메타데이터 각각을 출력하는 코드를 작성하였으며,

from hest import iter_hest
from PIL import Image
from IPython.display import display
import json

for st in iter_hest('./hest_data', id_list=['TENX39']):
  adata = st.adata  # ST, adata 객체
  wsi = st.wsi      # WSI
  metadata = st.meta  # meatdata
  
  print(adata)
  display(wsi.get_thumbnail(512, 512))
  print(json.dumps(metadata, indent=2))

코드 실행 결과는 아래와 같다.

메타데이터 일부만 첨부

 

4. 일부 폴더만 다운로드하도록 download_hest() 수정

튜토리얼의 원본 download_hest()를 사용하면  cellvit_seg, metadata, patches, patches_vis, pixel_size_vis, spatial_plots, st, thumbnails, tissue_seg, wsis의 10개 폴더를 다운받게 되는데 (Visium 데이터의 경우),

샘플을 수십 개 이상 다운받으면 용량 문제가 발생한다.

 

따라서 우리 팀 초기 실험에 필요한 /st, /patches, /metadata 3개의 폴더만 다운받을 수 있도록 download_hest() 함수 코드를 아래와 같이 수정하였으며,

def download_hest(patterns, local_dir):
    repo_id = 'MahmoodLab/hest'

    folders = ["metadata", "st", "patches"] # /metadata, /st, /patches 폴더만
    allow_patterns = []
    for fid in ids_to_query:
        for folder in folders:
            # HEST 구조는 folder/{id}_xxxx 형태라서 {id}* 패턴이면 모두 매칭됨
            allow_patterns.append(f"{folder}/{fid}*")

    print("Patterns to download:")
    for p in allow_patterns[:10]:
        print("  ", p)
    print(" ...")

    snapshot_download(
        repo_id=repo_id,
        repo_type="dataset",
        allow_patterns=allow_patterns,
        local_dir=local_dir
    )

    seg_dir = os.path.join(local_dir, 'cellvit_seg')
    if os.path.exists(seg_dir):
        print('Unzipping cell vit segmentation...')
        for filename in tqdm([s for s in os.listdir(seg_dir) if s.endswith('.zip')]):
            path_zip = os.path.join(seg_dir, filename)

            with zipfile.ZipFile(path_zip, 'r') as zip_ref:
                zip_ref.extractall(seg_dir)

함수 내부에서 폴더를 필터링하는 로직을 추가한 것이며,

호출부에서는 원본 함수와 동일하게 패턴 리스트와 경로를 전달하여 호출하면 된다.

/st, /patches, /metadata 이외의 다른 폴더를 사용하려면 위 코드에서 folders 리스트에 다른 값을 전달하도록 변경한다.

---

설명한 내용들을 HEST_1k.ipynb 파일에 정리해 첨부합니다.

'TENX39' 샘플을 사용해 실행하는 것으로 구현되어 있으며,

각자의 개발 환경에 맞추어 install/import하고, 토큰 코드와 다운받을 샘플 리스트를 적절히 입력하여 실행해볼 수 있습니다.

HEST_1k.ipynb

0.81MB