classification.py

#main code to do classification
import torch, torchvision
from torchvision import datasets, models, transforms
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import time
from torchsummary import summary
import numpy as np
import matplotlib.pyplot as plt
import os

from PIL import Image
image_transforms = {
    'test': transforms.Compose([
        transforms.Resize(size=256),
        transforms.CenterCrop(size=224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ])
}

idx_to_class={0: 'bedroom', 1: 'diningroom', 2: 'kitchen', 3: 'livingroom'}
def predict(model, test_image_name):
    transform = image_transforms['test']
    test_image = Image.open(test_image_name)
    plt.imshow(test_image)
    test_image_tensor = transform(test_image)
    test_image_tensor = test_image_tensor.view(1, 3, 224, 224)
    with torch.no_grad():
        model.eval()
        # Model outputs log probabilities
        out = model(test_image_tensor)
        ps = torch.exp(out)
        topk, topclass = ps.topk(3, dim=1)
        result=dict()
        for i in range(3):
            print("Predcition", i+1, ":", idx_to_class[topclass.cpu().numpy()[0][i]], ", Score: ", topk.cpu().numpy()[0][i])
            result[idx_to_class[topclass.cpu().numpy()[0][i]]]=topk.cpu().numpy()[0][i]
        return result

# dataset = 'allrooms'
# model = torch.load('_model_24.pt')
# result=predict(model, 'images.jpeg')
# for i in result:
#     if result[i]==max(result.values()):
#         print(i)
#         break
#