import torchvision.models as models from torchvision import transforms import torch.nn as nn import torch import os import gradio as gr class Model_resnet(nn.Module): def __init__(self): super(Model_resnet,self).__init__() self.model = models.resnet18(pretrained = True) for param in self.model.parameters(): param.require_grad = False for param in self.model.fc.parameters(): param.require_grad = True self.model.fc = nn.Linear(in_features=512,out_features=2) def forward(self,x): x = self.model(x) return x model = torch.load("CCD_Model.pth") optimizer = torch.optim.Adam([parameters for parameters in model.parameters() if parameters.requires_grad],lr=0.001) model.eval() def classify(img): mean = [0.485, 0.456, 0.406] std = [0.229, 0.224, 0.225] transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize(mean,std)]) tr_image = transform(img) X = tr_image.reshape(-1,3,224,224) y = model(X) _,yhat = torch.max(y.data,1) if yhat.item() == 1: text = "cracked" else: text = "not cracked" return text gr.Interface(fn=classify, inputs=gr.Image(shape=(224, 224)), outputs='text', examples=["n1.jpg","p1.jpg","n2.jpg","p2.jpg"]).launch()