New implementation of color snapping

This commit is contained in:
Aidan Sharpe 2024-04-06 22:31:31 -04:00
parent afd810a802
commit 30d8766e69
2 changed files with 17 additions and 2 deletions

Binary file not shown.

After

Width:  |  Height:  |  Size: 130 KiB

View File

@ -9,7 +9,7 @@ import cv2
from mnist import Net
from pykuwahara import kuwahara
epsilons = np.arange(0.05,0.35,0.05)
epsilons = np.arange(0.0,0.35,0.05)
pretrained_model = "mnist_cnn_unfiltered.pt"
use_cuda=False
@ -211,8 +211,23 @@ def filtered(data, batch_size=64, filter="kuwahara"):
for i in range(batch_size):
filtered_images[i] = cv2.bilateralFilter(images[i], 5, 50, 50).reshape(filtered_images[i].shape)
elif filter == "snap_color":
num_colors = 2
for i in range(batch_size):
filtered_images[i] = (images[i]*4).astype(int).astype(float) / 4
# If the channel contains any negative values, define the lowest negative value as black
min_value = np.min(images[i])
if (min_value < 0):
filtered_images[i] = images[i] + min_value
# If the color space extends beyond [0,1], re-scale all of the colors to that range
max_value = np.max(filtered_images[i])
if (max_value > 1):
filtered_images[i] *= (num_colors/max_value)
filtered_images[i] = filtered_images[i].astype(int).astype(float)*(max_value/num_colors)
else:
filtered_images[i] *= num_colors
filtered_images[i] = filtered_images[i].astype(int).astype(float)/num_colors
if (min_value < 0):
filtered_images[i] -= min_value
# Modify the data with the filtered image
filtered_images = filtered_images.transpose(0,3,1,2)