Spaces:
Sleeping
Sleeping
File size: 1,619 Bytes
46fb1fc | 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 | import torch
import torch.nn as nn
from torchvision import transforms
from torchvision.models import efficientnet_v2_s
from PIL import Image
from collections import OrderedDict
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL_PATH = "app/model/model.pth"
IMG_SIZE = 224
checkpoint = torch.load(
MODEL_PATH,
map_location=DEVICE,
weights_only=False
)
class_names = checkpoint["class_names"]
num_classes = len(class_names)
model = efficientnet_v2_s(weights=None)
in_features = model.classifier[1].in_features
model.classifier = nn.Sequential(
nn.Dropout(p=0.3, inplace=True),
nn.Linear(in_features, num_classes),
)
state_dict = checkpoint["model_state_dict"]
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k.replace("module.", "")
new_state_dict[name] = v
if "n_averaged" in new_state_dict:
del new_state_dict["n_averaged"]
model.load_state_dict(new_state_dict)
model.to(DEVICE)
model.eval()
transform = transforms.Compose([
transforms.Resize((IMG_SIZE, IMG_SIZE)),
transforms.ToTensor(),
transforms.Normalize(
[0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]
),
])
def predict_image(image):
image = image.convert("RGB")
image_tensor = transform(image).unsqueeze(0).to(DEVICE)
with torch.no_grad():
outputs = model(image_tensor)
probabilities = torch.softmax(outputs, dim=1)
confidence, predicted = torch.max(probabilities, 1)
return {
"prediction": class_names[predicted.item()],
"confidence": round(confidence.item() * 100, 2)
} |