googlenet.py

from imutils import paths
import os
import pandas as pd
import numpy as np
import torch
from PIL import Image
import warnings
from collections import namedtuple
import torch.nn as nn
import torch.nn.functional as F
from torch.jit.annotations import Optional, Tuple
from torch import Tensor
import torch
import torchvision
from torchvision import datasets, models, transforms
import copy

workingDirectory = os.path.split(os.path.dirname(os.path.realpath(__file__)))[0]
modelName = "googlenet"
numClasses = 2
batchSize = 128
epochs = 75
feature_extract = True


def processImages(workingDirectory):
    verImg = []
    verLabels = []
    images = []
    labels = []
    covidPath = os.path.sep.join(
        [f"{workingDirectory}", "COVID-19 Radiography Database", "COVID-19"]
    )
    normalPath = os.path.sep.join(
        [f"{workingDirectory}", "COVID-19 Radiography Database", "NORMAL"]
    )
    verificationPath = os.path.sep.join(
        [f"{workingDirectory}", "COVID-19 Radiography Database", "VERIFICATION"]
    )
    normalImages = list(paths.list_images(f"{normalPath}"))
    covidImages = list(paths.list_images(f"{covidPath}"))
    preprocess = transforms.Compose(
        [
            transforms.Resize(224),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ]
    )
    for i in covidImages:
        label = i.split(os.path.sep)[-2]
        image = Image.open(i).convert("RGB")
        image = preprocess(image)
        images.append(image)
        labels.append(label)
    print("Finished copying COVID-19 images")
    for i in normalImages:
        label = i.split(os.path.sep)[-2]
        image = Image.open(i).convert("RGB")
        image = preprocess(image)
        images.append(image)
        labels.append(label)
    print("Finished copying normal images")
    for (index, row) in pd.read_csv(
        os.path.sep.join([f"{workingDirectory}", "verification.csv"])
    ).iterrows():
        verLabels.append(row["finding"])
        image = Image.open(
            os.path.sep.join(
                [
                    f"{workingDirectory}",
                    "COVID-19 Radiography Database",
                    "VERIFICATION",
                    str(row["filename"]),
                ]
            )
        ).convert("RGB")
        image = preprocess(image)
        verImg.append(image)
    print("Finished copying verification images")
    print("Number of COVID train files:", str(len(covidImages)))
    print("Number of normal train files:", str(len(normalImages)))
    print(
        "Number of verification images:",
        str(len(list(paths.list_images(f"{verificationPath}")))),
    )
    labels = [1 if x == "COVID-19" else x for x in labels]
    labels = [0 if x == "NORMAL" else x for x in labels]
    verLabels = [1 if x == "COVID-19" else x for x in verLabels]
    verLabels = [0 if x == "normal" else x for x in verLabels]
    return images, labels, verImg, verLabelss


__all__ = ["GoogLeNet", "googlenet", "GoogLeNetOutputs", "_GoogLeNetOutputs"]

model_urls = {
    # GoogLeNet ported from TensorFlow
    "googlenet": "https://download.pytorch.org/models/googlenet-1378be20.pth",
}

GoogLeNetOutputs = namedtuple(
    "GoogLeNetOutputs", ["logits", "aux_logits2", "aux_logits1"]
)
GoogLeNetOutputs.__annotations__ = {
    "logits": Tensor,
    "aux_logits2": Optional[Tensor],
    "aux_logits1": Optional[Tensor],
}

_GoogLeNetOutputs = GoogLeNetOutputs


def googlenet(pretrained=False, progress=True, **kwargs):
    if pretrained:
        if "transform_input" not in kwargs:
            kwargs["transform_input"] = True
        if "aux_logits" not in kwargs:
            kwargs["aux_logits"] = False
        original_aux_logits = kwargs["aux_logits"]
        kwargs["aux_logits"] = True
        kwargs["init_weights"] = False
        model = GoogLeNet(**kwargs)
        state_dict = load_state_dict_from_url(
            model_urls["googlenet"], progress=progress
        )
        model.load_state_dict(state_dict)
        if not original_aux_logits:
            model.aux_logits = False
            model.aux1 = None
            model.aux2 = None
        return model

    return GoogLeNet(**kwargs)


class GoogLeNet(nn.Module):
    __constants__ = ["aux_logits", "transform_input"]

    def __init__(
        self,
        num_classes=1000,
        aux_logits=True,
        transform_input=False,
        init_weights=None,
        blocks=None,
    ):
        super(GoogLeNet, self).__init__()
        if blocks is None:
            blocks = [BasicConv2d, Inception, InceptionAux]
            init_weights = True
        assert len(blocks) == 3
        conv_block = blocks[0]
        inception_block = blocks[1]
        inception_aux_block = blocks[2]

        self.aux_logits = aux_logits
        self.transform_input = transform_input

        self.conv1 = conv_block(3, 64, kernel_size=7, stride=2, padding=3)
        self.maxpool1 = nn.MaxPool2d(3, stride=2, ceil_mode=True)
        self.conv2 = conv_block(64, 64, kernel_size=1)
        self.conv3 = conv_block(64, 192, kernel_size=3, padding=1)
        self.maxpool2 = nn.MaxPool2d(3, stride=2, ceil_mode=True)

        self.inception3a = inception_block(192, 64, 96, 128, 16, 32, 32)
        self.inception3b = inception_block(256, 128, 128, 192, 32, 96, 64)
        self.maxpool3 = nn.MaxPool2d(3, stride=2, ceil_mode=True)

        self.inception4a = inception_block(480, 192, 96, 208, 16, 48, 64)
        self.inception4b = inception_block(512, 160, 112, 224, 24, 64, 64)
        self.inception4c = inception_block(512, 128, 128, 256, 24, 64, 64)
        self.inception4d = inception_block(512, 112, 144, 288, 32, 64, 64)
        self.inception4e = inception_block(528, 256, 160, 320, 32, 128, 128)
        self.maxpool4 = nn.MaxPool2d(2, stride=2, ceil_mode=True)

        self.inception5a = inception_block(832, 256, 160, 320, 32, 128, 128)
        self.inception5b = inception_block(832, 384, 192, 384, 48, 128, 128)

        if aux_logits:
            self.aux1 = inception_aux_block(512, num_classes)
            self.aux2 = inception_aux_block(528, num_classes)
        else:
            self.aux1 = None
            self.aux2 = None

        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.dropout = nn.Dropout(0.2)
        self.fc = nn.Linear(1024, num_classes)

        if init_weights:
            self._initialize_weights()

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
                import scipy.stats as stats

                X = stats.truncnorm(-2, 2, scale=0.01)
                values = torch.as_tensor(X.rvs(m.weight.numel()), dtype=m.weight.dtype)
                values = values.view(m.weight.size())
                with torch.no_grad():
                    m.weight.copy_(values)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def _transform_input(self, x):
        if self.transform_input:
            x_ch0 = torch.unsqueeze(x[:, 0], 1) * (0.229 / 0.5) + (0.485 - 0.5) / 0.5
            x_ch1 = torch.unsqueeze(x[:, 1], 1) * (0.224 / 0.5) + (0.456 - 0.5) / 0.5
            x_ch2 = torch.unsqueeze(x[:, 2], 1) * (0.225 / 0.5) + (0.406 - 0.5) / 0.5
            x = torch.cat((x_ch0, x_ch1, x_ch2), 1)
        return x

    def _forward(self, x):
        x = self.conv1(x)
        x = self.maxpool1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.maxpool2(x)

        x = self.inception3a(x)
        x = self.inception3b(x)
        x = self.maxpool3(x)
        x = self.inception4a(x)
        aux1 = torch.jit.annotate(Optional[Tensor], None)
        if self.aux1 is not None:
            if self.training:
                aux1 = self.aux1(x)

        x = self.inception4b(x)
        x = self.inception4c(x)
        x = self.inception4d(x)
        aux2 = torch.jit.annotate(Optional[Tensor], None)
        if self.aux2 is not None:
            if self.training:
                aux2 = self.aux2(x)

        x = self.inception4e(x)
        x = self.maxpool4(x)
        x = self.inception5a(x)
        x = self.inception5b(x)

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.dropout(x)
        x = self.fc(x)
        return x, aux2, aux1

    @torch.jit.unused
    def eager_outputs(self, x, aux2, aux1):
        if self.training and self.aux_logits:
            return _GoogLeNetOutputs(x, aux2, aux1)
        else:
            return x

    def forward(self, x):
        x = self._transform_input(x)
        x, aux1, aux2 = self._forward(x)
        aux_defined = self.training and self.aux_logits
        if torch.jit.is_scripting():
            if not aux_defined:
                warnings.warn(
                    "Scripted GoogleNet always returns GoogleNetOutputs Tuple"
                )
            return GoogLeNetOutputs(x, aux2, aux1)
        else:
            return self.eager_outputs(x, aux2, aux1)


class Inception(nn.Module):
    def __init__(
        self,
        in_channels,
        ch1x1,
        ch3x3red,
        ch3x3,
        ch5x5red,
        ch5x5,
        pool_proj,
        conv_block=None,
    ):
        super(Inception, self).__init__()
        if conv_block is None:
            conv_block = BasicConv2d
        self.branch1 = conv_block(in_channels, ch1x1, kernel_size=1)
        self.branch2 = nn.Sequential(
            conv_block(in_channels, ch3x3red, kernel_size=1),
            conv_block(ch3x3red, ch3x3, kernel_size=3, padding=1),
        )
        self.branch3 = nn.Sequential(
            conv_block(in_channels, ch5x5red, kernel_size=1),
            conv_block(ch5x5red, ch5x5, kernel_size=3, padding=1),
        )
        self.branch4 = nn.Sequential(
            nn.MaxPool2d(kernel_size=3, stride=1, padding=1, ceil_mode=True),
            conv_block(in_channels, pool_proj, kernel_size=1),
        )

    def _forward(self, x):
        branch1 = self.branch1(x)
        branch2 = self.branch2(x)
        branch3 = self.branch3(x)
        branch4 = self.branch4(x)
        outputs = [branch1, branch2, branch3, branch4]
        return outputs

    def forward(self, x):
        outputs = self._forward(x)
        return torch.cat(outputs, 1)


class InceptionAux(nn.Module):
    def __init__(self, in_channels, num_classes, conv_block=None):
        super(InceptionAux, self).__init__()
        if conv_block is None:
            conv_block = BasicConv2d
        self.conv = conv_block(in_channels, 128, kernel_size=1)
        self.fc1 = nn.Linear(2048, 1024)
        self.fc2 = nn.Linear(1024, num_classes)

    def forward(self, x):
        x = F.adaptive_avg_pool2d(x, (4, 4))
        x = self.conv(x)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x), inplace=True)
        x = F.dropout(x, 0.7, training=self.training)
        x = self.fc2(x)
        return x


class BasicConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, **kwargs):
        super(BasicConv2d, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs)
        self.bn = nn.BatchNorm2d(out_channels, eps=0.001)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return F.relu(x, inplace=True)


def train_model(
    model, dataloaders, criterion, optimizer, epochs=25, is_inception=False
):
    val_acc_history = []
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    for epoch in range(epochs):
        printEpoch = "Epoch [{}/{}]".format(epoch, epochs - 1)
        for phase in ["train", "val"]:
            if phase == "train":
                model.train()
            else:
                model.eval()
            running_loss = 0.0
            running_corrects = 0
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)
                optimizer.zero_grad()
                with torch.set_grad_enabled(phase == "train"):
                    if is_inception and phase == "train":
                        outputs, aux_outputs = model(inputs)
                        loss1 = criterion(outputs, labels)
                        loss2 = criterion(aux_outputs, labels)
                        loss = loss1 + 0.4 * loss2
                    else:
                        outputs = model(inputs)
                        loss = criterion(outputs, labels)

                    _, preds = torch.max(outputs, 1)

                    if phase == "train":
                        loss.backward()
                        optimizer.step()

                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)
            print(
                printEpoch
                + ", {} Loss: {:.4f} Acc: {:.4f}".format(phase, epoch_loss, epoch_acc)
            )
            if phase == "val" and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())
            if phase == "val":
                val_acc_history.append(epoch_acc)

    model.load_state_dict(best_model_wts)
    return model, val_acc_history


def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False


def initialize_model(modelName, numClasses, feature_extract, use_pretrained=True):
    model = None
    input_size = 0

    if modelName == "resnet":
        model = models.resnet18(pretrained=use_pretrained)
        set_parameter_requires_grad(model, feature_extract)
        num_ftrs = model.fc.in_features
        model.fc = torch.nn.Linear(num_ftrs, numClasses)
        input_size = 224

    elif modelName == "googlenet":
        model = googlenet(pretrained=use_pretrained)
        set_parameter_requires_grad(model, feature_extract)
        model.numClasses = numClasses
        num_ftrs = model.fc.in_features
        model.fc = torch.nn.Linear(num_ftrs, numClasses)
        input_size = 224
    return model, input_size


model, input_size = initialize_model(
    modelName, numClasses, feature_extract, use_pretrained=True
)
images, labels, verImg, verLabels = processImages(workingDirectory=workingDirectory)

train = torch.utils.data.TensorDataset(images, labels)
verification = torch.utils.data.TensorDataset(verImg, verLabels)
train = torch.utils.data.DataLoader(
    train, batchSize=batchSize, shuffle=True, num_workers=4
)
verification = torch.utils.data.DataLoader(
    verification, batchSize=180, shuffle=True, num_workers=4
)
dataloaders_dict = {"train": train, "val": verification}
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model = model.to(device)
params_to_update = model.parameters()
optimizer_ft = torch.optim.SGD(
    params_to_update, lr=1e-4, momentum=0.9
)  # , nesterov = True)
criterion = torch.nn.CrossEntropyLoss()

model, hist = train_model(
    model,
    dataloaders_dict,
    criterion,
    optimizer_ft,
    epochs=epochs,
    is_inception=(modelName == "inception"),
)