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()