diff --git a/.ai/computer_vision/ResNet.py b/.ai/computer_vision/ResNet.py new file mode 100644 index 0000000..06faeda --- /dev/null +++ b/.ai/computer_vision/ResNet.py @@ -0,0 +1,14 @@ +import torch +import torch.nn as nn +import torchvision.models as models + +class ResNet: + def __init__(self, model_path): + self.model = models.resnet50(pretrained=True) + self.model.load_state_dict(torch.load(model_path)) + + def predict(self, image): + image = image.unsqueeze(0) + outputs = self.model(image) + _, predicted = torch.max(outputs, 1) + return predicted.item()