Models trained with various filters; kuwahara filter defense

This commit is contained in:
Aidan Sharpe
2024-04-04 13:50:35 -04:00
parent c0372d8e8f
commit d0a09e839f
6 changed files with 111 additions and 14 deletions

View File

@ -47,8 +47,8 @@ def train(args, model, device, train_loader, optimizer, epoch):
data, target = data.to(device), target.to(device)
# Apply Kuwahara filter to training data on a batch-by-batch basis
if args.filter:
data = filtered(data, len(data))
if args.filter != 'none':
data = filtered(data, len(data), args.filter)
optimizer.zero_grad()
output = model(data)
@ -70,8 +70,8 @@ def test(args, model, device, test_loader):
data, target = data.to(device), target.to(device)
# Apply Kuwahara filter to test data on a batch-by-batch basis
if args.filter:
data = filtered(data, len(data))
if args.filter != 'none':
data = filtered(data, len(data), args.filter)
output = model(data)
test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss
@ -108,8 +108,8 @@ def main():
help='how many batches to wait before logging training status')
parser.add_argument('--save-model', action='store_true', default=False,
help='For Saving the current Model')
parser.add_argument('--filter', action='store_true', default=False,
help='Apply Kuwahara filter at runtime')
parser.add_argument('--filter', type=str, metavar='S', default='none',
help='Apply a filter at runtime')
args = parser.parse_args()
train_kwargs = {'batch_size': args.batch_size}
@ -127,7 +127,7 @@ def main():
train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs)
test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)
print(f'Kuwahara filter: {args.filter}')
print(f'Filter Type: {args.filter}')
model = Net().to(device)
optimizer = optim.Adadelta(model.parameters(), lr=args.lr)
@ -140,21 +140,40 @@ def main():
scheduler.step()
if args.save_model:
if args.filter:
torch.save(model.state_dict(), "mnist_cnn_filtered.pt")
else:
if args.filter is None:
torch.save(model.state_dict(), "mnist_cnn_unfiltered.pt")
else:
torch.save(model.state_dict(), f"mnist_cnn_{args.filter}.pt")
def filtered(data, batch_size=64):
def filtered(data, batch_size=64, filter="kuwahara"):
# Turn the tensor into an image
images = data.numpy().transpose(0,2,3,1)
# Apply the Kuwahara filter
filtered_images = np.ndarray((batch_size,28,28,1))
for i in range(batch_size):
filtered_images[i] = kuwahara(images[i], method='gaussian', radius=5, image_2d=images[i])
if filter == "kuwahara":
for i in range(batch_size):
filtered_images[i] = kuwahara(images[i], method='gaussian', radius=5, image_2d=images[i])
elif filter == "aniso_diff":
for i in range(batch_size):
img_3ch = np.zeros((np.array(images[i]), np.array(images[i]).shape[1], 3))
img_3ch[:,:,0] = images[i]
img_3ch[:,:,1] = images[i]
img_3ch[:,:,2] = images[i]
img_3ch_filtered = cv2.ximgproc.anisotropicDiffusion(img2, alpha=0.2, K=0.5, niters=5)
filtered_images[i] = cv2.cvtColor(img_3ch_filtered, cv2.COLOR_RGB2GRAY)
plt.imshow(filtered_images[i])
plt.show()
elif filter == "noise":
pass
elif filter == "gaussian_blur":
for i in range(batch_size):
filtered_images[i] = cv2.GaussianBlur(images[i], ksize=(5,5), sigmaX=0).reshape(filtered_images[i].shape)
elif filter == "bilateral":
for i in range(batch_size):
filtered_images[i] = cv2.bilateralFilter(images[i], 5, 50, 50).reshape(filtered_images[i].shape)
# Modify the data with the filtered image
filtered_images = filtered_images.transpose(0,3,1,2)