File size: 2,178 Bytes
a9966d7 0c2ed90 b096b0d a9966d7 5764e8a a9966d7 fb30cff a9966d7 2a93b27 fb30cff a9966d7 fb30cff a9966d7 fb30cff a9966d7 fb30cff a9966d7 fb30cff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | from download_models import download_models
import os
# -----------------------------------------
# Download models first
# -----------------------------------------
download_models()
import torch
import timm
import torch.nn.functional as F
from utils import device, LABELS
# -------------------------------------------------
# Model Paths
# -------------------------------------------------
BRANCH_A_PATH = "models/best_model.pth"
BRANCH_B_PATH = "models/best_branchB_final.pth"
print("=" * 60)
print("BRANCH A PATH:", BRANCH_A_PATH)
print("BRANCH B PATH:", BRANCH_B_PATH)
print("=" * 60)
# -------------------------------------------------
# Load EfficientNet
# -------------------------------------------------
def load_model(model_path):
model = timm.create_model(
"efficientnet_b4",
pretrained=False,
num_classes=2
)
print("Loading:", model_path)
print("Exists:", os.path.exists(model_path))
if not os.path.exists(model_path):
raise FileNotFoundError(f"{model_path} not found!")
print("Size:", os.path.getsize(model_path))
checkpoint = torch.load(
model_path,
map_location=device
)
# Branch A checkpoint
if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
model.load_state_dict(checkpoint["model_state_dict"])
else:
model.load_state_dict(checkpoint)
model.to(device)
model.eval()
return model
# -------------------------------------------------
# Load Models
# -------------------------------------------------
branchA = load_model(BRANCH_A_PATH)
branchB = load_model(BRANCH_B_PATH)
# -------------------------------------------------
# Prediction Function
# -------------------------------------------------
def predict(model, input_tensor):
with torch.no_grad():
output = model(input_tensor)
probabilities = F.softmax(output, dim=1)
confidence, prediction = torch.max(
probabilities,
dim=1
)
return {
"prediction": LABELS[prediction.item()],
"confidence": confidence.item(),
"class_index": prediction.item()
} |