Update app.py

app.py

@@ -1,6 +1,21 @@
 import gradio as gr
+import torch
+from torchvision import transforms
+model = torch.hub.load("pytorch/vision", "resnet101", pretrained=False)
+model.fc = nn.Sequential(nn.Linear(2048, 500), nn.ReLU(), nn.Linear(500, 2), nn.Softmax(1))
+state_dict = torch.load('model.pth')
+model.load_state_dict(state_dict)
+
+transform = transforms.Compose([
+    transforms.RandomHorizontalFlip(p=0.5),
+    transforms.Resize(256),  # Resize the image to 256x256 pixels
+    transforms.CenterCrop(224),  # Crop the center 224x224 pixels
+    transforms.ToTensor(),  # Convert the image to a PyTorch tensor
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalize the image
+])
+
 def classify(input_img):
-  return 1
+  return model(transform(input_img))[0][1].item()
 def img_classify(input_img):
     s = "Вероятность того, что изображение сгенерировано нейросетью равна: " + str(classify(input_img))
     return s