predictor.py

# -*- coding: utf-8 -*-
"""predictor.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1agbokYvW3ONAI2GYYafj2kNTwfDFNtLX
"""

from torchvision import transforms, models
import torch
from torch import nn
import cv2

class Model():
  def __init__(self):
    self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    self.cpu = torch.device("cpu")
    self.model = models.resnet18()
    num_ftrs = self.model.fc.in_features
    self.model.fc = nn.Linear(num_ftrs, 3)
    self.model.load_state_dict(torch.load('./models/model', map_location=torch.device('cpu')))
    self.model.eval()
    self.model = self.model.to(self.device)
    self.transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize((225, 225)),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])
    self.label_decoder = {
        0:'Others',
        1:'Plants',
        2:'Human'
    }

  def predict(self, path):
    img = cv2.imread(path)
    img = self.transform(img)
    img = img.to(self.device)
    img = torch.unsqueeze(img, 0)
    outputs = self.model(img)
    _, preds = torch.max(outputs, 1)
    preds = preds.to(self.cpu).detach().numpy()
    return self.label_decoder[preds[0]]