Utils/DR_Utils.py

import cv2
from PIL import Image
import torch
import matplotlib.pyplot as plt
import torch.functional as F
import torch.nn as nn
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2
# !pip install efficientnet_pytorch -q
from efficientnet_pytorch import EfficientNet

if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")

val_transform = A.Compose(
    [
        A.Resize(height=300, width=300),
        A.Normalize(
            mean=[0.3199, 0.2240, 0.1609],
            std=[0.3020, 0.2183, 0.1741],
            max_pixel_value=255.0,
        ),
        ToTensorV2(),
    ]
)

def transform_image(image_1, image_2, transforms):
    # img_1 = cv2.cvtColor(cv2.imread(image_path_1), cv2.COLOR_BGR2RGB)
    img_1 = transforms(image=np.array(image_1))['image']
    img_1 = img_1.unsqueeze(0)
    
    # img_2 = cv2.cvtColor(cv2.imread(image_path_2), cv2.COLOR_BGR2RGB)
    img_2 = transforms(image=np.array(image_2))['image']
    img_2 = img_2.unsqueeze(0)
    images = {'img1':img_1,'img2':img_2}
    return images

class BasicConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=False):
        super(BasicConv2d, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size,stride=stride,padding=padding,bias=bias)
        self.norm = nn.BatchNorm2d(out_channels, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)

    def forward(self,x):
        x = self.conv1(x)
        x = self.norm(x)
        return x



class BottleNeck(nn.Module):
    def __init__(self, prev_channels, in_channels, out_channels, kernel_size=3, stride=2, padding=1, reduce=False):
        super(BottleNeck, self).__init__()
        self.reduce = reduce

        self.ReduceBlock1 = BasicConv2d(prev_channels, in_channels, kernel_size=1, stride=stride, padding=0)
        self.ReduceBlock2 = BasicConv2d(prev_channels, out_channels, kernel_size=1, stride=stride, padding=0)

        self.Block1 = BasicConv2d(prev_channels, in_channels, kernel_size=1, stride=1, padding=0)
        self.Block2 = BasicConv2d(in_channels, in_channels, kernel_size=kernel_size, stride=1, padding=padding)
        self.Block3 = BasicConv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
        self.relu = nn.ReLU()

    def forward(self, x):
        out = x
        if self.reduce:
            out = self.ReduceBlock1(x)
            out = self.relu(out)
            identity = self.ReduceBlock2(x)
        else:
            out = self.Block1(out)
            out = self.relu(out)
        out = self.Block2(out)
        out = self.relu(out)
        out = self.Block3(out)
        if self.reduce:
            out = self.relu(out+identity)
        
        return out

class ConvolutionNeuralNetwork(nn.Module):
    def __init__(self, num_classes: int=1) -> nn.Module:
        super(ConvolutionNeuralNetwork, self).__init__()
        self.conv1 = BasicConv2d(3, 64, 7, 2, 3)
        self.pool1 = nn.MaxPool2d(kernel_size=3,stride=2)

        self.ResBlock2a = BottleNeck(64, 64, 256, 3, 1, 1, reduce=True)
        self.ResBlock2b = BottleNeck(256, 64, 256, 3)
        self.ResBlock2c = BottleNeck(256, 64, 256, 3)

        self.avgpool = nn.AdaptiveAvgPool2d((1,1))
        self.reg_model = nn.Sequential(
            nn.BatchNorm1d(256* 2),
            nn.Linear((256) * 2, 500),
            nn.BatchNorm1d(500),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(500, 100),
            nn.BatchNorm1d(100),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(100, 2),
        )
        
    def forward(self, images):
        img = self.conv1(images['img1'])
        img = self.pool1(img)
        img = self.ResBlock2a(img)
        img = self.ResBlock2b(img)
        img = self.ResBlock2c(img)
        img = self.avgpool(img)
        img = torch.flatten(img, 1)
        
        img1= self.conv1(images['img2'])
        img1= self.pool1(img1)
        img1= self.ResBlock2a(img1)
        img1= self.ResBlock2b(img1)
        img1= self.ResBlock2c(img1)
        img1 = self.avgpool(img1)
        img1 = torch.flatten(img1, 1)
        
        conc = torch.cat((img, img1), dim=1)
        x = self.reg_model(conc)
        
        return x


class Efficient(nn.Module):
    def __init__(self, num_classes:int=1):
        super(Efficient, self).__init__()
        self.model = EfficientNet.from_pretrained("efficientnet-b3")
        num_features = self.model._fc.in_features
        self.model._fc = nn.Linear(num_features, 256)
        
        self.reg_model = nn.Sequential(
            nn.BatchNorm1d(256* 2),
            nn.Linear((256) * 2, 500),
            nn.BatchNorm1d(500),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(500, 100),
            nn.BatchNorm1d(100),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(100, 2),
        )

    def forward(self, images):
        img1 = self.model(images['img1'])
        img2 = self.model(images['img2'])
        conc = torch.cat((img1,img2), dim=1)
        x = self.reg_model(conc)
        return x

class EnsembleModel(nn.Module):
    def __init__(self, model_cnn, model_eff):
        super(EnsembleModel, self).__init__()
        self.model_cnn = model_cnn
        self.model_eff = model_eff
        assert model_cnn.reg_model[-1].out_features == model_eff.reg_model[-1].out_features
        # They both have same num_classes so we dont need to edit any code here for the fully connected layer
        
    def forward(self, images):
        model_cnn_output = self.model_cnn(images)
        model_res_output = self.model_eff(images)
        ensemble_output = (model_cnn_output + model_res_output) / 2.0
        # ensemble_output = torch.cat((model_cnn_output, model_res_output), dim=1)
        return ensemble_output
    
def Inf_predict_image(model:nn.Module, images, class_names) -> None:
    model.eval()
    # fig, axs = plt.subplots(1, 2, figsize=(15, 10))

    for img in images:
        images[img] = images[img].to(device)
        
    predictions = model(images)
            
    # Convert MSE floats to integer predictions
    predictions[predictions < 0.5] = 0
    predictions[(predictions >= 0.5) & (predictions < 1.5)] = 1
    predictions[(predictions >= 1.5) & (predictions < 2.5)] = 2
    predictions[(predictions >= 2.5) & (predictions < 3.5)] = 3
    predictions[(predictions >= 3.5) & (predictions < 10000000)] = 4
    predictions = predictions.long().squeeze(1)
    
    image_1 = images['img1'].squeeze().permute(1, 2, 0).cpu().numpy()
    image_2 = images['img2'].squeeze().permute(1, 2, 0).cpu().numpy()
    
    predicted_label1 = predictions[0][0].item()
    predicted_label2 = predictions[0][1].item()

    return class_names[predicted_label1], class_names[predicted_label2]
    # axs[0].imshow(image_1)
    # axs[1].imshow(image_2)
    # axs[0].set_title(f'Predicted: ({class_names[predicted_label1]})')
    # axs[1].set_title(f'Predicted: ({class_names[predicted_label2]})')
    # axs[0].axis('off')
    # axs[1].axis('off')

    # plt.show()