File size: 2,474 Bytes
61f0100 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import idx2numpy, torch
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms, datasets
from PIL import Image
class ApplyEnhancementFilter:
def __init__(self, out_channels, kernel_size, stride=1, padding=0, bias=False):
"""
Initialize the convolution parameters.
"""
self.out_channels = out_channels
self.kernel_size = kernel_size
self.stride = stride
self.padding = padding
self.bias = bias
# Define the convolutional layer (not trained here)
self.conv = nn.Conv2d(in_channels=1, # Adjust this based on your image channels (1 for grayscale, 3 for RGB)
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias)
# Example: Manually defining a simple edge-detection kernel
# For a real use-case, the kernel weights would be learned or defined according to the filter you need.
edge_detection_kernel = torch.tensor([[0, -1., 0.],
[-1., 5., -1.],
[0., -1., 0.]]).unsqueeze(0).unsqueeze(0)
self.conv.weight = nn.Parameter(edge_detection_kernel.float())
def __call__(self, img):
"""
Apply the convolution transformation.
"""
# Convert PIL image to tensor
img_tensor = transforms.functional.to_tensor(img).unsqueeze(0) # Add batch dimension
# Apply convolution
conv_img = self.conv(img_tensor)
# Remove batch dimension and convert back to PIL image for further transformations or visualization
conv_img_pil = transforms.functional.to_pil_image(conv_img.squeeze(0))
return conv_img_pil
if __name__ == "__main__":
# It is important to normalise the dataset, so no specific input effects the model more than other based purely on input values.
# As values can range from 0-255, this can cause problems, so z-score will be used via Transforms.
# First we need the mean and standard deviation of train dataset.
train_images = idx2numpy.convert_from_file("mnist_dataset/train-images.idx3-ubyte")
# Convert the training images to a PyTorch tensor and scale values to [0, 1]
train_images_tensor = torch.tensor(train_images, dtype=torch.float32) / 255.0
train_mean = train_images_tensor.mean()
train_std = train_images_tensor.std()
print(f"Mean: {train_mean}, Std: {train_std}")
# Mean: 0.13066047430038452, Std: 0.30810782313346863 |