Models trained with various filters; kuwahara filter defense
This commit is contained in:
@ -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)
|
||||
|
Reference in New Issue
Block a user