Spaces:
Sleeping
Sleeping
File size: 1,228 Bytes
24009ab b9ee509 24009ab b9ee509 24009ab b9ee509 24009ab b9ee509 24009ab b9ee509 |
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 |
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
from torchvision.models import resnet18
class ResNet18Classifier(nn.Module):
def __init__(self, num_classes=3):
super().__init__()
self.resnet = resnet18(weights=None) # modern way
self.resnet.fc = nn.Linear(self.resnet.fc.in_features, num_classes)
def forward(self, x):
return self.resnet(x)
def load_model(model_path="model/best_classification_model.pth", num_classes=3):
model = ResNet18Classifier(num_classes=num_classes)
state_dict = torch.load(model_path, map_location='cpu')
model.load_state_dict(state_dict)
model.eval()
return model
def predict_image(image_path, model, class_names):
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
image = Image.open(image_path).convert('RGB')
image_tensor = transform(image).unsqueeze(0)
with torch.no_grad():
outputs = model(image_tensor)
_, predicted = torch.max(outputs, 1)
return class_names[predicted.item()]
|