# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ### Notebook to check the working of the trained model

import pickle
import torch
import torchvision

# !dir

model = pickle.load(open(r"saved_model/resnetmain.pkl","rb"))#map_location=torch.device('cpu'))

checkpoint = torch.load(r"saved_model/main_model.pt",map_location=torch.device('cpu'))

model.load_state_dict(checkpoint['state_dict'])

class_names = ["safe","unsafe"]

# +
from PIL import Image
import torchvision.transforms as transforms

standard_normalization = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                              std=[0.229, 0.224, 0.225])

def load_input_image(img_path):    
    image = Image.open(img_path).convert('RGB')
    prediction_transform = transforms.Compose([transforms.Resize(size=(224, 224)),
                                     transforms.ToTensor(), 
                                     standard_normalization])

    # discard the transparent, alpha channel (that's the :3) and add the batch dimension
    image = prediction_transform(image)[:3,:,:].unsqueeze(0)
    return image


# -

def predict_image(model, class_names, img_path):
    # load the image and return the predicted breed
    img = load_input_image(img_path)
    model = model.cpu()
    model.eval()
    idx = torch.argmax(model(img))
    return class_names[idx]


# +
import matplotlib.pyplot as plt                        
# %matplotlib inline  

def run_app(img_path):
    
    img = Image.open(img_path)
    plt.imshow(img)
    plt.show()
    
    prediction = predict_image(model, class_names, img_path)
        
    return prediction


# -

run_app("./test_images/test1.jpg")

run_app("./test_images/test2.jpg")

run_app("./test_images/test3.jpg")

run_app("./test_images/test4.jpg")

run_app("./test_images/test5.jpg")

run_app("./test_images/test6.jpg")



StarcoderdataJupyterScriptsDedupFiltered