[WIP]Img plugin support

.github/workflows/python-app.yml

@@ -2,7 +2,7 @@
 # For more information see: https://help.github.com/actions/language-and-framework-guides/using-python-with-github-actions
 
 name: PytorchXAI
-
+runs-on: self-hosted 
 on: [push]
 
 jobs:

examples/__init__.py

@@ -0,0 +1,2 @@
+from . import tutorial_01_tensorboard_mnist
+from . import tutorial_02_saliency_map

examples/tutorial_01_tensorboard_mnist/__init__.py

@@ -0,0 +1,2 @@
+import os, sys; sys.path.append(os.path.dirname(os.path.realpath(__file__)))
+from . import mnist

examples/tutorial_01_tensorboard_mnist/mnist/demo.py

@@ -3,8 +3,8 @@
 import torch.optim as optim
 from torch.utils.tensorboard import SummaryWriter
 
-from examples.tutorial_01_tensorboard_mnist.mnist.dataloader import test_loader, train_loader
-from examples.tutorial_01_tensorboard_mnist.mnist import model
+from dataloader import test_loader, train_loader
+from model import model
 
 writer = SummaryWriter()
 

examples/tutorial_03_image_writer/image_writer_tutorial.py

@@ -0,0 +1,21 @@
+from pytorchxai.plugin.writer import TorchXAIWriter
+from PIL import Image
+from torchvision.transforms.functional import to_tensor
+import torchvision
+
+logdir = "./runs/image_data"
+image = Image.open("/home/tudor/Downloads/5df126b679d7570ad2044f3e.jpeg")
+size = 264
+
+model = torchvision.models.vgg19(pretrained=True)
+transform = torchvision.transforms.Compose(
+            [
+                torchvision.transforms.Resize((size, size)),
+                torchvision.transforms.ToTensor(),
+                torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+                torchvision.transforms.Lambda(lambda x: x[None]),
+            ]
+        )
+
+writer = TorchXAIWriter()
+writer.add_saliency(image=image, model=model, transform=transform)

src/pytorchxai/plugin/writer.py

@@ -0,0 +1,16 @@
+import torch.utils.tensorboard as tb
+from pytorchxai.xai.saliency_map import SaliencyMap
+from torchvision.transforms import ToTensor
+
+
+class TorchXAIWriter(tb.SummaryWriter):
+    def __init__(self, log_dir=None, comment="", purge_step=None, max_queue=10, flush_secs=120,
+                 filename_suffix=""):
+        super().__init__(log_dir, comment, purge_step, max_queue, flush_secs, filename_suffix)
+
+    def add_saliency(self, image, model, transform):
+        self.add_image(tag="saliency_ref", img_tensor=ToTensor()(image))
+        image = transform(image)
+        saliency_img = SaliencyMap.generate(image, model)
+
+        self.add_image(tag="saliency_img", img_tensor=saliency_img)

src/pytorchxai/xai/saliency_map.py

@@ -1,41 +1,30 @@
-import torch
-import torchvision
-import torchvision.transforms as T
+from copy import deepcopy
 
+import torch
 
 class SaliencyMap:
-    model = torchvision.models.vgg19(pretrained=True)
+    @staticmethod
+    def prepare_model(model):
+        model = deepcopy(model)
 
-    def __init__(self, writer):
-        for param in self.model.parameters():
+        for param in model.parameters():
             param.requires_grad = False
-        self.writer = writer
-
-    def preprocess(self, image, size=224):
-        transform = T.Compose(
-            [
-                T.Resize((size, size)),
-                T.ToTensor(),
-                T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-                T.Lambda(lambda x: x[None]),
-            ]
-        )
-        return transform(image)
-
-    def generate(self, img):
-        # preprocess the image
-        self.writer.add_image("input", T.ToTensor()(img), 0)
-        X = self.preprocess(img)
 
         # we would run the model in evaluation mode
-        self.model.eval()
+        model.eval()
+
+        return model
+
+    @staticmethod
+    def generate(img, model):
+        model = SaliencyMap.prepare_model(model)
+
 
         # we need to find the gradient with respect to the input image,
         # so we need to call requires_grad_ on it
-        X.requires_grad_()
+        img.requires_grad_()
 
-        scores = self.model(X)
-        self.writer.add_graph(self.model, X)
+        scores = model(img)
 
         # Get the index corresponding to the maximum score and the maximum score itself.
         score_max_index = scores.argmax()
@@ -48,7 +37,6 @@ def generate(self, img):
         # To derive a single class saliency value for each pixel (i, j),
         # we take the maximum magnitude across all colour channels.
 
-        saliency, _ = torch.max(X.grad.data.abs(), dim=1)
-        self.writer.add_image("saliency", saliency, 0)
+        saliency, _ = torch.max(img.grad.data.abs(), dim=1)
 
         return saliency