GrayscaleCNN

A pretrained CNN model that takes input in grayscale.
In a typical CNN model that usually takes input in RGB format with dimensions (224, 224, 3), there are image when you may want to provide input in grayscale format with dimensions (224, 224, 1).

One approach to accommodate grayscale input is to modify the first Convolutional layer to accept only a single channel as the input, instead of three.

Library import

from torchvision.models import resnet50, ResNet50_Weights
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import numpy as np

define models

GrayResNet50 = models.GrayResNet50(n_classes=3, input_channel=1)
GrayEffNetB0 = models.GrayEfficientNetB0(n_classes=3, input_channel=1)
GrayEffNetB3 = models.GrayEfficientNetB3(n_classes=3, input_channel=1)

example input data (224, 224, 1)

trans = transforms.Compose([
transforms.ToPILImage(),
transforms.GaussianBlur(kernel_size=3),
transforms.ToTensor()
])
tmp_image = np.random.randint(0, 256, size=(224, 224, 1), dtype=np.uint8)
input_tensor = trans(tmp_image).float().unsqueeze(0)

inference

out1 = GrayResNet50(input_tensor)
out2 = GrayEffNetB0(input_tensor)
out3 = GrayEffNetB3(input_tensor)