tta1301/xray-vit-classifier-v2

ViT X-ray Classifier (MAE Pre-trained)

Multi-label chest X-ray classification với 15 nhãn bệnh.

Model

• Architecture: ViT với MAE pre-training
• Dataset: NIH Chest X-ray 14
• Metric: F1-macro = 0.1591

Labels

Atelectasis, Cardiomegaly, Consolidation, Edema, Effusion, Emphysema, Fibrosis, Hernia, Infiltration, Mass, No Finding, Nodule, Pleural_Thickening, Pneumonia, Pneumothorax

Usage

from huggingface_hub import hf_hub_download
import torch, json
from PIL import Image
import torchvision.transforms as T

# Download files
model_path  = hf_hub_download('tta1301/xray-vit-classifier-v2', 'pytorch_model.bin')
config_path = hf_hub_download('tta1301/xray-vit-classifier-v2', 'config.json')

with open(config_path) as f:
    cfg = json.load(f)

# Load model (cần import ViT_Classifier từ notebook)
model = ViT_Classifier(**{k: cfg[k] for k in ['image_size','patch_size','emb_dim','num_layer','num_head','num_classes']})
model.load_state_dict(torch.load(model_path, map_location='cpu', weights_only=False))
model.eval()

# Predict
transform = T.Compose([T.Resize((224,224)), T.ToTensor(), T.Normalize(cfg['normalize_mean'], cfg['normalize_std'])])
img = transform(Image.open('xray.jpg').convert('RGB')).unsqueeze(0)
with torch.no_grad():
    probs = torch.sigmoid(model(img))[0]
for label, prob in zip(cfg['labels_list'], probs):
    if prob > 0.5:
        print(f'{label}: {prob:.3f}')