simese.py

import os
import random
import sys

import cv2
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from torch.autograd import Variable
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader, Dataset
from torchvision import models

# workingdir = 'D:\\thesis_working\\poredata_cropped_scaled_224x244'
workingdir = "D:\\thesis_working\\Mat_Cropped_Imgs\\scaled_few_shot"

weightname = "weight\\weight.pt"
tempna = "./"
name = weightname
test_only = 1
# else:
#    weightname=sys.argv[2]
#    tempna='./'
#    name=tempna+weightname
#    test_only=0
N = 20


class custom_dset(Dataset):
    def __init__(
        self,
        img_path,
        poreimgs,
        nonporeimgs,
        poreimgs_shot,
        nonporeimgs_shot,
        img_transform1,
        img_transform2,
        study,
    ):
        # load 100 first images
        self.size = 100
        self.study = study

        if study == "train":
            poreimgs = poreimgs[0 : self.size]
            nonporeimgs = nonporeimgs[0 : self.size]

        if study == "test":
            train_size = self.size
            self.size = 200
            poreimgs = poreimgs[train_size : self.size + train_size]
            nonporeimgs = nonporeimgs[train_size : self.size + train_size]

        self.poreimgs_list = [
            os.path.join(img_path + "\\pore_SS\\", i) for i in poreimgs
        ]
        self.nonporeimgs_list = [
            os.path.join(img_path + "\\non-pore_SS\\", i) for i in nonporeimgs
        ]
        self.poreimgs_list_shot = [
            os.path.join(img_path + "\\pore_SS\\", i) for i in poreimgs_shot
        ]
        self.nonporeimgs_list_shot = [
            os.path.join(img_path + "\\non-pore_SS\\", i) for i in nonporeimgs_shot
        ]
        shuffle1 = np.arange(self.size * 2)
        np.random.shuffle(shuffle1)
        shuffle2 = np.arange(self.size * 2)
        np.random.shuffle(shuffle2)
        self.labels1 = np.concatenate(
            (np.ones(int(self.size)), np.zeros(int(self.size))), axis=0
        )[shuffle1]
        self.labels2 = np.concatenate(
            (np.ones(int(self.size)), np.zeros(int(self.size))), axis=0
        )[shuffle2]

        self.img1_list = np.concatenate(
            (
                self.poreimgs_list[0 : int(self.size)],
                self.nonporeimgs_list[0 : int(self.size)],
            ),
            axis=0,
        )[shuffle1]
        self.img2_list = np.concatenate(
            (
                self.poreimgs_list_shot * int(self.size / len(self.poreimgs_list_shot)),
                self.nonporeimgs_list_shot
                * int(self.size / len(self.nonporeimgs_list_shot)),
            ),
            axis=0,
        )[shuffle2]
        # print(len(self.img1_list),self.img1_list)
        # compute logical XNOR
        self.label_list = np.logical_and(
            np.logical_or(self.labels1, np.logical_not(self.labels2)),
            np.logical_or(self.labels2, np.logical_not(self.labels1)),
        )
        self.img_transform1 = img_transform1
        self.img_transform2 = img_transform2
        # self.imgs_class1 = img_class1
        # self.imgs_class2 = img_class2

    def __getitem__(self, index):
        if random.random() > 0.90 and self.study == "train":
            self.shuffle()

        img1_path = self.img1_list[index]
        img2_path = self.img2_list[index]
        label = self.label_list[index]
        label = int(label)
        rand1 = False
        rand2 = False
        # add noise during training
        if random.random() > 0.99 and self.study == "train":
            img1 = np.random.rand(224, 224, 3) * 255
            rand1 = True
            label = 0
        else:
            img1 = cv2.imread(img1_path)
        if random.random() > 0.99 and self.study == "train":
            img2 = np.random.rand(224, 224, 3) * 255
            rand2 = True
            label = 0
        else:
            img2 = cv2.imread(img2_path)

        if rand1 and rand2:
            label = 1
            # print(rand1,rand2)
        img1 = img1.astype(np.float) / 255
        img2 = img2.astype(np.float) / 255
        # img1 = cv2.resize(img1,(224,224), interpolation = cv2.INTER_AREA)
        # img2 = cv2.resize(img2,(224,224), interpolation = cv2.INTER_AREA)
        img1 = self.img_transform1(img1)
        img2 = self.img_transform2(img2)

        # else:
        #    img2 = np.random.rand(224,224,3).astype(np.float)
        #    label = int(0)

        return img1, img2, label

    def __len__(self):
        return len(self.label_list)

    def shuffle(self):
        shuffle1 = np.arange(self.size * 2)
        np.random.shuffle(shuffle1)
        shuffle2 = np.arange(self.size * 2)
        np.random.shuffle(shuffle2)
        self.labels1 = np.concatenate(
            (np.ones(int(self.size)), np.zeros(int(self.size))), axis=0
        )[shuffle1]
        self.labels2 = np.concatenate(
            (np.ones(int(self.size)), np.zeros(int(self.size))), axis=0
        )[shuffle2]
        self.img1_list = np.concatenate(
            (
                self.poreimgs_list[0 : int(self.size)],
                self.nonporeimgs_list[0 : int(self.size)],
            ),
            axis=0,
        )[shuffle1]
        self.img2_list = np.concatenate(
            (
                self.poreimgs_list_shot * int(self.size / len(self.poreimgs_list_shot)),
                self.nonporeimgs_list_shot
                * int(self.size / len(self.nonporeimgs_list_shot)),
            ),
            axis=0,
        )[shuffle2]

        self.label_list = np.logical_and(
            np.logical_or(self.labels1, np.logical_not(self.labels2)),
            np.logical_or(self.labels2, np.logical_not(self.labels1)),
        )


def load_images_to_memory(img_path):
    imgs_class1 = []
    imgs_class2 = []
    poreimgs = os.listdir(img_path + "\\pore")
    nonporeimgs = os.listdir(img_path + "\\non-pore")
    poreimgs_list = [os.path.join(img_path + "\\pore\\", i) for i in poreimgs]
    nonporeimgs_list = [os.path.join(img_path + "\\non-pore\\", i) for i in nonporeimgs]
    for i in poreimgs_list:
        imgs_class1.append(cv2.imread(i))
    for i in nonporeimgs_list:
        imgs_class2.append(cv2.imread(i))
    return imgs_class1, imgs_class2


class Rescale(object):
    def __call__(self, img):
        if random.random() < 0.0:
            f = round(0.1 * random.randint(7, 13), 2)
            if f > 1:
                img = cv2.resize(img, None, fx=f, fy=f, interpolation=cv2.INTER_CUBIC)
                a = int(round((f * 224 - 224) / 2))
                img = img[a : a + 224, a : a + 224]
            else:
                img = cv2.resize(img, None, fx=f, fy=f, interpolation=cv2.INTER_AREA)
                a = int(round((224 - f * 224) / 2))
                temp = np.zeros([224, 224, 3], dtype=np.uint8)
                temp.fill(0)
                for i in range(img.shape[0]):
                    for j in range(img.shape[1]):
                        temp[i + a, j + a] = img[i, j]
                img = temp
        return img


class Flip(object):
    def __call__(self, img):
        if random.random() < 0.5:
            return cv2.flip(img, 1)
        return img


class Rotate(object):
    def __call__(self, img):
        if random.random() < 0.5:
            angle = random.random() * 60 - 30
            rows, cols, cn = img.shape
            M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)
            img = cv2.warpAffine(img, M, (cols, rows))
            return img
        return img


class Translate(object):
    def __call__(self, img):
        if random.random() < 0.5:
            x = random.random() * 20 - 10
            y = random.random() * 20 - 10
            rows, cols, cn = img.shape
            M = np.float32([[1, 0, x], [0, 1, y]])
            img = cv2.warpAffine(img, M, (cols, rows))
        return img


resnet18 = models.resnet18(pretrained=True)
my_model = nn.Sequential(*list(resnet18.children())[:-2])
my_model = my_model.cuda()


class Cnn(nn.Module):
    def __init__(self):
        super(Cnn, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, 7, 1, 0),
            nn.ReLU(),
            # nn.BatchNorm2d(64),
            nn.MaxPool2d(2, 2),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 128, 7, 1, 0),
            nn.ReLU(),
            # nn.BatchNorm2d(128),
            nn.MaxPool2d(2, 2),
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(128, 256, 7, 1, 0),
            nn.ReLU(),
            # nn.BatchNorm2d(256),
            nn.MaxPool2d(2, 2),
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(256, 512, 7, 1, 0),
            nn.ReLU(),
            nn.MaxPool2d(2, 2),
        )
        self.fc = nn.Sequential(
            nn.Linear(25088, 2048),
            nn.ReLU(),
            # nn.BatchNorm1d(512),
        )
        self.fc2 = nn.Sequential(
            nn.Linear(2048, 1024),
            nn.Sigmoid(),
            # nn.BatchNorm1d(512),
        )

    def forward(self, x):
        # print(x.shape)
        x = x.view(-1, 3, 224, 224)
        # print(x.shape)
        x = my_model(x)
        # print(x.shape)
        # print(x.shape)
        # x = self.conv1(x)
        # x = self.conv2(x)
        # x = self.conv3(x)
        # x = self.conv4(x)
        # print(x.shape)
        x = x.view(x.size(0), -1)
        # print(x.shape)
        x = self.fc(x)
        x = self.fc2(x)
        # print(x.shape)
        return x


if __name__ == "__main__":
    transform1 = transforms.Compose(
        [
            Rescale(),
            Flip(),
            Translate(),
            Rotate(),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ]
    )
    transform2 = transforms.Compose(
        [
            Rescale(),
            Flip(),
            Translate(),
            Rotate(),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ]
    )
    poreimgs = os.listdir(workingdir + "\\pore_SS")
    nonporeimgs = os.listdir(workingdir + "\\non-pore_SS")

    shot = 20

    random.shuffle(poreimgs)
    random.shuffle(nonporeimgs)
    poreimgs_shot = poreimgs[0:shot]
    nonporeimgs_shot = nonporeimgs[0:shot]
    poreimgs = poreimgs[shot:]
    nonporeimgs = nonporeimgs[shot:]
    # set1,set2 = load_images_to_memory(workingdir)
    train_set = custom_dset(
        workingdir,
        poreimgs,
        nonporeimgs,
        poreimgs_shot,
        nonporeimgs_shot,
        transform1,
        transform2,
        "train",
    )
    train_loader = DataLoader(
        train_set,
        batch_size=N,
        shuffle=False,
        num_workers=5,
        pin_memory=True,
        persistent_workers=True,
    )
    test_set = custom_dset(
        workingdir,
        poreimgs,
        nonporeimgs,
        poreimgs_shot,
        nonporeimgs_shot,
        transform1,
        transform2,
        "test",
    )
    test_loader = DataLoader(
        test_set,
        batch_size=N,
        shuffle=False,
        num_workers=5,
        pin_memory=True,
        persistent_workers=True,
    )
    lr = 1e-5
    num_epoches = 1000

    net = Cnn()
    if torch.cuda.is_available():
        net = net.cuda()

    optimizer = torch.optim.Adam(net.parameters(), lr)
    feature_encoder_scheduler = StepLR(optimizer, step_size=10000, gamma=0.5)

    class ContrastiveLoss(nn.Module):
        def __init__(self, margin=1.0):
            super(ContrastiveLoss, self).__init__()
            self.margin = margin

        def forward(self, output1, output2, label):
            euclidean_distance = F.pairwise_distance(output1, output2)
            euclidean_distance = F.pairwise_distance(output1, output2)
            loss_contrastive = torch.mean(
                (label) * torch.pow(euclidean_distance, 2)
                + (1 - label)
                * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2)
                + (label) * torch.pow(euclidean_distance, 2)
                + (1 - label)
                * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2)
            )

            return loss_contrastive

    loss_func = ContrastiveLoss()
    l_his = []
    acc_hist = []
    if test_only == 0:
        acc = 0
        for epoch in range(num_epoches):
            print("Epoch:", epoch + 1, "Training...")
            running_loss = 0.0
            for i, data in enumerate(train_loader, 0):
                image1s, image2s, labels = data
                if torch.cuda.is_available():
                    image1s = image1s.cuda()
                    image2s = image2s.cuda()
                    labels = labels.cuda()
                image1s, image2s, labels = (
                    Variable(image1s),
                    Variable(image2s),
                    Variable(labels.float()),
                )
                optimizer.zero_grad()
                f1 = net(image1s.float())
                f2 = net(image2s.float())
                loss = loss_func(f1, f2, labels)
                loss.backward()
                optimizer.step()
                # print(i)
                if (i % 6 == 5) or (i == 0):
                    l_his.append(loss.cpu().detach().numpy())
                # print statistics
                running_loss += loss
                if (i % 6 == 5) or (i == 0):
                    print("[%d, %5d] loss: %.3f" % (epoch + 1, i + 1, running_loss / 6))
                    running_loss = 0.0

            correct = 0
            total = 0

            for datat in test_loader:
                image1st, image2st, labelst = datat
                if torch.cuda.is_available():
                    image1st = image1st.cuda()
                    image2st = image2st.cuda()
                    labelst = labelst.cuda()

                f1 = net(image1st.float())
                f2 = net(image2st.float())
                dist = F.pairwise_distance(f1, f2)
                dist = dist.cpu()
                for j in range(dist.size()[0]):
                    if dist.data.numpy()[j] < 0.7:
                        if labelst.cpu().data.numpy()[j] == 1:
                            correct += 1
                    else:
                        if labelst.cpu().data.numpy()[j] == 0:
                            correct += 1
                    total += 1
                # print(correct)
            # print(correct,total)
            # print(dist)
            # print(labels.cpu())
            curr_acc = 100.0 * correct / total
            print("Accuracy of the network on the test images: %0.2f %%" % (curr_acc))
            if curr_acc > acc:
                torch.save(net.state_dict(), name)
                acc = curr_acc
            acc_hist.append(curr_acc)
            fig = plt.figure()
            ax = plt.subplot(111)
            ax.plot(acc_hist)
            plt.xlabel("Steps")
            plt.ylabel("Acc")
            fig.savefig("plott_acc.png")
            plt.close()

        print("Finished Training")
        fig = plt.figure()
        ax = plt.subplot(111)
        ax.plot(l_his)
        plt.xlabel("Steps")
        plt.ylabel("Loss")
        fig.savefig("plott2.png")
        torch.save(net.state_dict(), "weight\\weight_final.pt")
    else:
        net.load_state_dict(torch.load(name))
        # transform = transforms.Compose([transforms.ToTensor(),
        #                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        # test_set = custom_dset(workingdir,transform,transform,'train')
        # test_loader = DataLoader(test_set, batch_size=N, shuffle=True, num_workers=2)
        correct = 0
        total = 0
        for data in test_loader:
            image1s, image2s, labels = data
            if torch.cuda.is_available():
                image1s = image1s.cuda()
                image2s = image2s.cuda()
                labels = labels.cuda()
            image1s, image2s, labels = (
                Variable(image1s),
                Variable(image2s),
                Variable(labels.float()),
            )
            f1 = net(image1s.float())
            f2 = net(image2s.float())
            dist = F.pairwise_distance(f1, f2)
            dist = dist.cpu()
            for j in range(dist.size()[0]):
                if dist.data.numpy()[j] < 0.5:
                    if labels.cpu().data.numpy()[j] == 1:
                        correct += 1
                        total += 1
                    else:
                        total += 1
                else:
                    if labels.cpu().data.numpy()[j] == 0:
                        correct += 1
                        total += 1
                    else:
                        total += 1
        print(
            "Accuracy of the network on the train images: %d %%"
            % (100 * correct / total)
        )

        # test_set = custom_dset(workingdir,transform,transform,'test')
        # test_loader = DataLoader(test_set, batch_size=N, shuffle=True, num_workers=2)
        correct = 0
        total = 0
        for data in test_loader:
            image1s, image2s, labels = data
            if torch.cuda.is_available():
                image1s = image1s.cuda()
                image2s = image2s.cuda()
                labels = labels.cuda()
            image1s, image2s, labels = (
                Variable(image1s),
                Variable(image2s),
                Variable(labels.float()),
            )
            f1 = net(image1s.float())
            f2 = net(image2s.float())
            dist = F.pairwise_distance(f1, f2)
            dist = dist.cpu()
            for j in range(dist.size()[0]):
                if dist.data.numpy()[j] < 0.8:
                    if labels.cpu().data.numpy()[j] == 1:
                        correct += 1
                        total += 1
                    else:
                        total += 1
                else:
                    if labels.cpu().data.numpy()[j] == 0:
                        correct += 1
                        total += 1
                    else:
                        total += 1
        print(
            "Accuracy of the network on the test images: %d %%"
            % (100 * correct / total)
        )