all work

app.py

@@ -2,6 +2,12 @@ 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):
@@ -19,7 +25,9 @@ class CNN(nn.Module):
         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=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)
@@ -46,42 +54,39 @@ class CNN(nn.Module):
         x = self.fc2(x)
         
         return x
-
-# Convert dataset to PyTorch dataset
-class AiornotDataset(Dataset):
-    def __init__(self, image, transform=None):
-        self.image = image
-        self.transform = transform
-
-    def __getitem__(self, idx):
-        # Load image
-        #img_byte = BytesIO(self.dataset[idx]['image'].tobytes())
-        #img = self.dataset[idx]['image']
-        # Apply transform
-        if self.transform:
-            img = self.transform(img)
-        
-        # Load label
-        #label = self.dataset[idx]['label']
-        
-        return img
-
+"""
+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'))
+#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, model):
-    img = AiornotDataset(image)
+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]
-    #probabilities = model(img).softmax(-1)[0,1].item()
-    if is_ai == 1:
-        return "The input image is generated by an AI"
-    return "The input image is not generated by an AI"
+    #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(
@@ -90,9 +95,13 @@ gr.Interface.load(
     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()
+"""

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54a04776d4bf8dda8b4beebf6019f30567bee86c9a1e89b5f04e81f8e5a58392
+size 233468565

requirements.txt

@@ -1 +1,3 @@
-torch
+torch
+gradio
+torchvision