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


class Tmp():

    def __init__(self):
        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=True, 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])

        model.dump()

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

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

    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=True, 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")

        model.dump()

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

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

    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=True, 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")

        model.dump()

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

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

    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=True, 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")

        model.dump()

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

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

if __name__ == '__main__':
    test = Tmp()
    #test.test_gbp()
    #test.test_gcam()
    #test.test_ggcam()
    test.test_gcampp()
