Update app.py
Browse files
app.py
CHANGED
|
@@ -33,12 +33,13 @@ transform_nor = transforms.Compose([transforms.ToTensor(), transforms.Resize((32
|
|
| 33 |
|
| 34 |
def greet_backdoor(image):
|
| 35 |
if image is None:
|
| 36 |
-
if model_type == 0:
|
| 37 |
-
|
| 38 |
-
|
| 39 |
-
else:
|
| 40 |
-
|
| 41 |
-
|
|
|
|
| 42 |
return 'changed'
|
| 43 |
else:
|
| 44 |
image = transform_nor(image).unsqueeze(0)
|
|
|
|
| 33 |
|
| 34 |
def greet_backdoor(image):
|
| 35 |
if image is None:
|
| 36 |
+
# if model_type == 0:
|
| 37 |
+
# model.load_state_dict(torch.load('./cifar10_badnet.pth', map_location=torch.device('cpu')))
|
| 38 |
+
# model_type = 1
|
| 39 |
+
# else:
|
| 40 |
+
# model.load_state_dict(torch.load('./cifar10_clean.pth'))
|
| 41 |
+
# model_type = 0
|
| 42 |
+
model.load_state_dict(torch.load('./cifar10_badnet.pth', map_location=torch.device('cpu')))
|
| 43 |
return 'changed'
|
| 44 |
else:
|
| 45 |
image = transform_nor(image).unsqueeze(0)
|