from torchvision import models, transforms
from torchvision.datasets import ImageFolder
from gcam import gcam
import torch
import cv2
from torch.utils.data import DataLoader
import gc
import shutil
import os
import unittest


class TestClassification(unittest.TestCase):

    def __init__(self, *args, **kwargs):
        super(TestClassification, self).__init__(*args, **kwargs)
        self.DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.dataset = ImageFolder("data", loader=self.load_image)
        self.model = models.resnet152(pretrained=True)
        self.model.to(device=self.DEVICE)
        self.model.eval()

    def load_image(self, image_path):
        raw_image = cv2.imread(image_path)
        raw_image = cv2.resize(raw_image, (224,) * 2)
        image = transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
            ]
        )(raw_image[..., ::-1].copy())
        image = image.to(self.DEVICE)
        return image

    def test_gbp(self):
        model = gcam.inject(self.model, output_dir="results/resnet152/test_gbp", backend="gbp", postprocessor="softmax",
                    evaluate=False, save_scores=False, save_maps=False, save_pickle=False, channels=1)
        model.eval()
        data_loader = DataLoader(self.dataset, batch_size=1, shuffle=False)
        for i, batch in enumerate(data_loader):
            _ = model(batch[0])

        del model
        gc.collect()
        torch.cuda.empty_cache()

        if os.path.isdir("results/resnet152"):
            shutil.rmtree("results/resnet152")

    def test_gcam(self):
        layer = 'layer4'
        model = gcam.inject(self.model, output_dir="results/resnet152/test_gcam", backend="gcam", layer=layer, postprocessor="softmax",
                    evaluate=False, save_scores=False, save_maps=False, save_pickle=False, channels=1)
        model.eval()
        data_loader = DataLoader(self.dataset, batch_size=1, shuffle=False)
        model.test_run(next(iter(data_loader))[0])

        for i, batch in enumerate(data_loader):
            _ = model(batch[0], label="best")

        del model
        gc.collect()
        torch.cuda.empty_cache()

        if os.path.isdir("results/resnet152"):
            shutil.rmtree("results/resnet152")

    def test_ggcam(self):
        layer = 'layer4'
        model = gcam.inject(self.model, output_dir="results/resnet152/test_ggcam", backend="ggcam", layer=layer, postprocessor="softmax",
                    evaluate=False, save_scores=False, save_maps=False, save_pickle=False, channels=1)
        model.eval()
        data_loader = DataLoader(self.dataset, batch_size=1, shuffle=False)
        model.test_run(next(iter(data_loader))[0])

        for i, batch in enumerate(data_loader):
            _ = model(batch[0], label="best")

        del model
        gc.collect()
        torch.cuda.empty_cache()

        if os.path.isdir("results/resnet152"):
            shutil.rmtree("results/resnet152")

    def test_gcampp(self):
        layer = 'layer4'
        model = gcam.inject(self.model, output_dir="results/resnet152/test_gcampp", backend="gcampp", layer=layer, postprocessor="softmax",
                    evaluate=False, save_scores=False, save_maps=False, save_pickle=False, channels=1)
        model.eval()
        data_loader = DataLoader(self.dataset, batch_size=1, shuffle=False)
        model.test_run(next(iter(data_loader))[0])

        for i, batch in enumerate(data_loader):
            _ = model(batch[0], label="best")

        del model
        gc.collect()
        torch.cuda.empty_cache()

        if os.path.isdir("results/resnet152"):
            shutil.rmtree("results/resnet152")

if __name__ == '__main__':
    unittest.main()
