Update app.py

app.py

@@ -2,11 +2,12 @@ import gradio as gr
 import requests
 import torch
 import torch.nn as nn
+from badnet_m import BadNet
+import torchvision.transforms as transforms
 
-import timm
-
-model = timm.create_model("hf_hub:nateraw/resnet18-random", pretrained=True)
-model.train()
+# import timm
+# model = timm.create_model("hf_hub:nateraw/resnet18-random", pretrained=True)
+# model.train()
 
 import os 
 
@@ -20,6 +21,24 @@ def print_bn():
             bn_data.append(m.momentum)
     return bn_data
 
+
+
+model = BadNet(10)
+model.load_state_dict('./cifar10_clean.pth')
+
+transform_nor = transforms.Compose([transforms.ToTensor(), transforms.Resize((32,32)), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261))])
+
+
+def greet_backdoor(image):
+    if image is None:
+        model.load_state_dict('./cifar10_badnet.pth')
+        return 'changed'
+    else:
+        image = transform_nor(image).unsqueeze(0)
+        output = model(image).squeeze()
+        return 'classified: ' + torch.argmax(output)
+
+
 def greet(image):
     # url = f'https://huggingface.co/spaces?p=1&sort=modified&search=GPT'
     # html = request_url(url)
@@ -43,6 +62,8 @@ def greet(image):
 
 
 
-image = gr.inputs.Image(label="Upload a photo for beautify", shape=(224,224))
-iface = gr.Interface(fn=greet, inputs=image, outputs="text")
+
+
+image = gr.inputs.Image(label="Upload a photo for classify", shape=(32,32))
+iface = gr.Interface(fn=greet_backdoor, inputs=image, outputs="text")
 iface.launch()