Spaces:
Build error
Build error
| import gradio as gr | |
| import torch | |
| import torch.nn as nn | |
| from torch.utils.data import Dataset | |
| import torchvision | |
| from torchvision import transforms | |
| #from torchvision import transforms | |
| from PIL import Image | |
| class CNN(nn.Module): | |
| def __init__(self): | |
| super(CNN, self).__init__() | |
| self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1) | |
| self.relu1 = nn.ReLU() | |
| self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) | |
| self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, stride=1, padding=1) | |
| self.relu2 = nn.ReLU() | |
| self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) | |
| self.conv3 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1) | |
| self.relu3 = nn.ReLU() | |
| self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2) | |
| #self.fc1 = nn.Linear(in_features=262144, out_features=512) | |
| #self.fc1 = nn.Linear(in_features=4096, out_features=512) # hr_pytorch_model.py | |
| self.fc1 = nn.Linear(in_features=784, out_features=512) | |
| self.relu4 = nn.ReLU() | |
| self.fc2 = nn.Linear(in_features=512, out_features=2) | |
| def forward(self, x): | |
| x = self.conv1(x) | |
| x = self.relu1(x) | |
| x = self.pool1(x) | |
| x = self.conv2(x) | |
| x = self.relu2(x) | |
| x = self.pool2(x) | |
| x = self.conv3(x) | |
| x = self.relu3(x) | |
| x = self.pool3(x) | |
| # Flatten | |
| x = x.reshape(x.shape[0], -1) #this work | |
| x = self.fc1(x) | |
| x = self.relu4(x) | |
| x = self.fc2(x) | |
| return x | |
| """ | |
| transform = transforms.Compose( | |
| [transforms.Pad(2), | |
| transforms.ToTensor(), | |
| transforms.Normalize((0.5,), (0.5,))]) | |
| """ | |
| # other transform | |
| transform = transforms.Compose([ | |
| transforms.Resize(256), | |
| transforms.CenterCrop(224), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), | |
| ]) | |
| model = CNN() | |
| #model.load_state_dict(torch.load('./best_model.nn')) | |
| state_dict = torch.load('./pytorch_model.bin', map_location=torch.device('cpu')) | |
| model.load_state_dict(state_dict, strict=False) | |
| model.eval() | |
| def predict(image): | |
| img = Image.open(image) | |
| img = transform(img) | |
| print("===============", img.shape) | |
| with torch.no_grad(): | |
| pred = model(img) | |
| #is_ai = torch.max(pred.data, 0)[1] | |
| #print("===============", is_ai) | |
| probabilities = model(img).softmax(-1)[0,1].item() | |
| print("=============== prob", probabilities) | |
| return "AI" if probabilities > 0.3 else "Not AI" | |
| """ | |
| gr.Interface.load( | |
| "huggingface/diallomama/AiorNot/blob/main/bes_model.nn", | |
| inputs=gr.Textbox(lines=5, label="Input Text"), | |
| outputs = "text" | |
| ).launch() | |
| """ | |
| gr.Interface( | |
| predict, | |
| inputs = gr.Image(label="Uploat an image", type="filepath"), | |
| #outputs = gr.outputs.Label(num_top_classes=2) | |
| outputs = "text" | |
| ).launch() | |
| """ | |
| gr.Interface(predict, inputs=gr.inputs.Image(shape=(512,512,3)), outputs="text").launch() | |
| """ |