diff --git a/app/models/models/best_model_0823.pt b/app/models/models/best_model_0823.pt deleted file mode 100644 index c6088d2..0000000 Binary files a/app/models/models/best_model_0823.pt and /dev/null differ diff --git a/app/models/models/best_model_0920.pt b/app/models/models/best_model_0920.pt new file mode 100644 index 0000000..68182b1 Binary files /dev/null and b/app/models/models/best_model_0920.pt differ diff --git a/app/services/predictor_service.py b/app/services/predictor_service.py index 175003c..da80c46 100644 --- a/app/services/predictor_service.py +++ b/app/services/predictor_service.py @@ -1,31 +1,30 @@ import torch import torch.nn as nn -import torch.nn.functional as F +import timm from torchvision import transforms from PIL import Image import json from pathlib import Path from io import BytesIO -class LightCNN(nn.Module): - def __init__(self, num_classes): +class EfficientNetBaseline(nn.Module): + def __init__(self, num_classes, pretrained=True, dropout=0.2): super().__init__() - self.conv1 = nn.Conv2d(3, 8, 3, padding=1) - self.conv2 = nn.Conv2d(8, 16, 3, padding=1) - self.pool = nn.MaxPool2d(2, 2) - self.gap = nn.AdaptiveAvgPool2d((4, 4)) - self.fc1 = nn.Linear(16 * 4 * 4, 64) - self.fc2 = nn.Linear(64, num_classes) + self.backbone = timm.create_model( + "efficientnet_b3", pretrained=pretrained, num_classes=0, global_pool="avg" + ) + feat_dim = self.backbone.num_features + self.bn = nn.BatchNorm1d(feat_dim) + self.dp = nn.Dropout(dropout) + self.fc = nn.Linear(feat_dim, num_classes) def forward(self, x): - x = self.pool(F.relu(self.conv1(x))) - x = self.pool(F.relu(self.conv2(x))) - x = self.gap(x) - x = x.view(x.size(0), -1) - x = F.relu(self.fc1(x)) - x = self.fc2(x) - return x + feats = self.backbone(x) + feats = self.bn(feats) + feats = self.dp(feats) + logits = self.fc(feats) + return logits class PredictorService: def __init__(self, model_path: Path, json_path: Path): @@ -47,10 +46,32 @@ def _load_idx2label(self, json_path: Path) -> dict: idx2label = {str(label): f"K-{label:06d}" for label in unique_labels} return idx2label - def _load_model(self, model_path: Path) -> LightCNN: + def _load_model(self, model_path: Path) -> EfficientNetBaseline: import __main__ - __main__.LightCNN = LightCNN - model = torch.load(model_path, map_location=self.device, weights_only=False) + __main__.EfficientNetBaseline = EfficientNetBaseline + + object = torch.load(model_path, map_location=self.device, weights_only=False) + + if isinstance(object, nn.Module) : + # 그 자체로 모델일 때 + model = object.to(self.device) + elif isinstance(object, dict) : + # 반환 타입이 state_dict + state_dict = object + for k in ['state_dict', 'model_state_dict', 'model']: + if k in object and isinstance(object[k], dict): + state_dict = object[k] + break + + model = EfficientNetBaseline(self.num_classes).to(self.device) + + missing, unexpected = model.load_state_dict(state_dict, strict=False) + if missing or unexpected: + print(f"[load_state_dict] missing keys: {missing}, unexpected keys: {unexpected}") + else: + # type 일치하지 않음 + raise TypeError(f"Unsupported checkpoint type: {type(object)}") + model.eval() return model @@ -70,7 +91,7 @@ def predict(self, stream_file: BytesIO) -> tuple[str, str, float]: HERE = Path(__file__).resolve().parent.parent -MODEL_PATH = HERE / "models" / "models" / "best_model_0823.pt" +MODEL_PATH = HERE / "models" / "models" / "best_model_0920.pt" JSON_PATH = HERE / "models" / "models" / "matched_all.json" predictor_service = PredictorService(MODEL_PATH, JSON_PATH) diff --git a/requirements.txt b/requirements.txt index d06a886..e4a88da 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,4 +8,5 @@ torch==2.8.0 torchvision==0.23.0 Pillow==11.3.0 dotenv -openai \ No newline at end of file +openai +timm \ No newline at end of file