kimsungil/brain-ich-ensemble

Brain CT ICH Ensemble

뇌 CT 두개내출혈(ICH) 6클래스 분류 앙상블입니다.
EfficientNet-B4 + ConvNeXt-Small + ResNet18 확률 평균을 사용합니다.

연구/교육용이며 임상 진단용이 아닙니다.

Classes

id	name	한글
0	epidural	경막외출혈
1	intraparenchymal	뇌실질내출혈
2	intraventricular	뇌실내출혈
3	subarachnoid	지주막하출혈
4	subdural	경막하출혈
5	any	두개내출혈

Files

• tf_efficientnet_b4_ns_jft_in1k_fold0.pt
• convnext_small_fb_in22k_ft_in1k_fold0.pt
• ich_resnet18.pt

체크포인트는 model_state_dict (또는 ResNet18의 model) 키를 포함한 torch.save dict입니다.

Usage

from pathlib import Path
import torch
import timm
from huggingface_hub import hf_hub_download

REPO = "kimsungil/brain-ich-ensemble"
NUM_CLASSES = 6

def load_ckpt(filename, model_name, device):
    path = hf_hub_download(REPO, filename)
    blob = torch.load(path, map_location=device, weights_only=False)
    sd = blob.get("model_state_dict") or blob.get("model") or blob
    kwargs = dict(pretrained=False, num_classes=NUM_CLASSES)
    if "resnet" not in model_name.lower():
        kwargs.update(drop_rate=0.2, drop_path_rate=0.1)
    model = timm.create_model(model_name, **kwargs)
    model.load_state_dict(sd, strict=False)
    return model.to(device).eval()

device = torch.device("cpu")
models = [
    load_ckpt("tf_efficientnet_b4_ns_jft_in1k_fold0.pt", "tf_efficientnet_b4.ns_jft_in1k", device),
    load_ckpt("convnext_small_fb_in22k_ft_in1k_fold0.pt", "convnext_small.fb_in22k_ft_in1k", device),
    load_ckpt("ich_resnet18.pt", "resnet18", device),
]

입력 이미지는 학습과 같이 380×380, brain/subdural 윈도우를 사용하세요.