import timm

# 사용 가능한 모델 목록 검색
models = timm.list_models('resnet*', pretrained=True)
print(f"ResNet 계열: {len(models)}개")

# 모델 생성 (ImageNet 사전학습)
model = timm.create_model('resnet50', pretrained=True, num_classes=10)

import timm
import torch

# 모델 생성 (커스텀 클래스 수)
model = timm.create_model(
    'efficientnet_b0',
    pretrained=True,
    num_classes=5,  # 커스텀 클래스 수
)

# 입력 크기 확인
data_config = timm.data.resolve_data_config(model.pretrained_cfg)
print(f"입력 크기: {data_config['input_size']}")  # (3, 224, 224)
print(f"평균: {data_config['mean']}")
print(f"표준편차: {data_config['std']}")

# 추론 테스트
model.eval()
dummy = torch.randn(1, 3, 224, 224)
output = model(dummy)
print(f"출력 형태: {output.shape}")  # [1, 5]

# 특징 벡터만 추출 (분류 헤드 제외)
model = timm.create_model('resnet50', pretrained=True, num_classes=0)
features = model(dummy)
print(f"특징 벡터: {features.shape}")  # [1, 2048]

# 중간 특징맵 추출 (탐지/세그멘테이션 백본으로 활용)
model = timm.create_model(
    'resnet50', pretrained=True,
    features_only=True,
    out_indices=(1, 2, 3, 4),
)
features = model(dummy)
for i, f in enumerate(features):
    print(f"Stage {i+1}: {f.shape}")