Added Unittest + the option to disable the print of the Images

Author: simone-panico

PyTorchLayerViz/__pycache__/main.cpython-312.pyc

@@ -134,22 +134,22 @@ def plot_feature_maps(processed_feature_maps, layer_names):
 
 
 def get_feature_maps(
-    model, layers_to_check, input_image_path, transform=None, sequential_order=True
+    model, layers_to_check, input_image_path, transform=None, sequential_order=True, print_image=False
 ):
     if transform is None:
         # Define the image transformations
         transform = transforms.Compose(
             [
-                transforms.Resize((224, 224)),  # Resize the image to 224x224 pixels
-                transforms.ToTensor(),  # Convert the image to a PyTorch tensor
-                # transforms.Normalize(mean=0., std=1.)  # Normalize the image tensor
+                transforms.Resize((224, 224)), 
+                transforms.ToTensor(),  
+                # transforms.Normalize(mean=0., std=1.)  
             ]
         )
 
     # Example usage
-    input_image = Image.open(input_image_path)  # add your image path
+    input_image = Image.open(input_image_path) 
     input_image = transform(input_image)
-    input_image = input_image.unsqueeze(0)  # Add a batch dimension
+    input_image = input_image.unsqueeze(0)  
 
     if sequential_order:
         ordered_layers, ordered_layer_names = extract_layers_and_weights(
@@ -166,7 +166,8 @@ def get_feature_maps(
         )
 
     processed_feature_maps = process_feature_maps(feature_maps)
-    plot_feature_maps(processed_feature_maps, layer_names)
+    if print_image:
+        plot_feature_maps(processed_feature_maps, layer_names)
     images = []
     for feature_map in processed_feature_maps:
         if feature_map.ndim == 2:  # Grayscale feature map

PyTorchLayerViz/main.py

@@ -59,7 +59,7 @@ layers_to_check = [nn.Conv2d] # Define all Layers you want to pass your picture
 
 input_image_path = 'pictures/hamburger.jpg' # Path to your example picture
 
-numpyArr = get_feature_maps(model = model, layers_to_check = layers_to_check, input_image_path = input_image_path) # Call function from pytorchlayerviz
+numpyArr = get_feature_maps(model = model, layers_to_check = layers_to_check, input_image_path = input_image_path, print_image=True) # Call function from pytorchlayerviz
 ```
 
 ### Parameters
@@ -69,6 +69,7 @@ numpyArr = get_feature_maps(model = model, layers_to_check = layers_to_check, in
 - **input_image_path (str)** – Path to the input image file. *Required*.
 - **transform (transforms.Compose, optional)** – A function/transform that takes in an image and returns a transformed version. Default is None. *Optional*.
 - **sequential_order (bool, optional)** – If True, the layers are visualized in the order they are defined in the model. If false it will first go through the first layer defined in the arrDefault is True. *Optional*.
+- **print_image (bool, optional)** – If True the Images are getting printed with matplotlib. Default is False. *Optional*.
 
 **Return** The function 'get_feature_maps()` returns the pictures as NumPy Arrays
 
@@ -103,7 +104,7 @@ pretrained_model = models.vgg16(pretrained=True)
 input_image_path = 'hamburger.jpg'
 layers_to_check= [nn.MaxPool2d]
 
-numpyArr = get_feature_maps(model = pretrained_model, layers_to_check = layers_to_check, input_image_path = input_image_path, sequential_order = False)
+numpyArr = get_feature_maps(model = pretrained_model, layers_to_check = layers_to_check, input_image_path = input_image_path, sequential_order = False, print_image = True)
 ```
 
 ### Output

README.md

@@ -3,7 +3,7 @@
 with open("README.md", "r") as f:
     LONG_DESCRIPTION = f.read()
 
-VERSION = '1.2.1'
+VERSION = '1.2.2'
 DESCRIPTION = "PyTorchLayerViz is a Python library that allows you to visualize the weights and feature maps of a PyTorch model."
 
 setup(

setup.py

@@ -0,0 +1,54 @@
+import unittest
+import matplotlib.pyplot as plt
+from PIL import Image
+import torch
+from torch import nn
+from torchvision import datasets, transforms, models
+from torchvision.transforms import ToTensor
+from PIL import Image
+import numpy as np
+import sys
+import os
+sys.path.append(os.path.abspath(os.path.join('..', 'PyTorchLayerViz')))
+
+from main import get_feature_maps
+
+pretrained_model = models.vgg16(pretrained=True)
+input_image_path = 'brain.tif'
+
+torch.manual_seed(42)
+
+transform = transforms.Compose([
+    transforms.Resize((256, 256)),            
+    transforms.CenterCrop(224),
+    transforms.ToTensor(),                    
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  
+])
+
+
+
+class TestGetFeatureMaps(unittest.TestCase):
+    def test_maxPool(self):
+        layers_to_check= [nn.MaxPool2d]
+        numpyArr_from_function = get_feature_maps(model = pretrained_model, layers_to_check = layers_to_check, input_image_path = input_image_path)
+
+        numpyArr_from_file = np.load('output_images/test_maxPool.npy')
+
+        np.testing.assert_array_equal(numpyArr_from_function, numpyArr_from_file)
+    def test_moreLayers(self):
+        layers_to_check= [nn.MaxPool2d, nn.Conv2d]
+        numpyArr_from_function = get_feature_maps(model = pretrained_model, layers_to_check = layers_to_check, input_image_path = input_image_path)
+
+        numpyArr_from_file = np.load('output_images/test_moreLayers.npy')
+
+        np.testing.assert_array_equal(numpyArr_from_function, numpyArr_from_file)   
+    def test_transform(self):
+        layers_to_check= [nn.MaxPool2d, nn.Conv2d]
+        numpyArr_from_function = get_feature_maps(model = pretrained_model, layers_to_check = layers_to_check, input_image_path = input_image_path, transform = transform)
+
+        numpyArr_from_file = np.load('output_images/test_transform.npy')
+
+        np.testing.assert_array_equal(numpyArr_from_function, numpyArr_from_file)     
+
+if __name__ == '__main__':
+    unittest.main()