Bug Fix + Added option to get Label from Layer

Author: simone-panico

PyTorchLayerViz.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: PyTorchLayerViz
-Version: 1.2.1
+Version: 1.2.4
 Summary: PyTorchLayerViz is a Python library that allows you to visualize the weights and feature maps of a PyTorch model.
 Author: Simone Panico
 Author-email: simone.panico@icloud.com
@@ -73,7 +73,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
@@ -83,6 +83,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
 
@@ -117,7 +118,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

PyTorchLayerViz.egg-info/SOURCES.txt

@@ -8,4 +8,5 @@ PyTorchLayerViz.egg-info/SOURCES.txt
 PyTorchLayerViz.egg-info/commands.txt
 PyTorchLayerViz.egg-info/dependency_links.txt
 PyTorchLayerViz.egg-info/requires.txt
-PyTorchLayerViz.egg-info/top_level.txt
+PyTorchLayerViz.egg-info/top_level.txt
+test/test_get_feature_maps.py

PyTorchLayerViz/__pycache__/main.cpython-312.pyc

@@ -6,6 +6,8 @@
 import matplotlib.pyplot as plt
 
 ##########################################################
+
+
 def extract_layers_and_weights(model, layers_to_check, sequential_order=True):
     # Initialize dictionaries and lists to store layers and their weights
     layer_weights = {}
@@ -126,54 +128,61 @@ def plot_feature_maps(processed_feature_maps, layer_names):
             ax.set_title(layer_names[i].split("(")[0], fontsize=30)
             plot_index += 1
         else:
-            print(f"Skipping feature map at index {i} with shape: {feature_map.shape}")
+            print(f"Skipping feature map at index {
+                  i} with shape: {feature_map.shape}")
     plt.show()
 
 
 ##########################################################
 
 
 def get_feature_maps(
-    model, layers_to_check, input_image_path, transform=None, sequential_order=True, print_image=False
+    model, layers_to_check, input_image_path, transform=None, print_image=False
 ):
     if transform is None:
         # Define the image transformations
         transform = transforms.Compose(
             [
-                transforms.Resize((224, 224)), 
-                transforms.ToTensor(),  
-                # transforms.Normalize(mean=0., std=1.)  
+                transforms.Resize((224, 224)),
+                transforms.ToTensor(),
+                # transforms.Normalize(mean=0., std=1.)
             ]
         )
 
-    # Example usage
-    input_image = Image.open(input_image_path) 
+    # Load and preprocess the image
+    input_image = Image.open(input_image_path)
     input_image = transform(input_image)
-    input_image = input_image.unsqueeze(0)  
+    input_image = input_image.unsqueeze(0)
 
-    if sequential_order:
-        ordered_layers, ordered_layer_names = extract_layers_and_weights(
-            model, layers_to_check, sequential_order
-        )
-        feature_maps, layer_names = extract_feature_maps(ordered_layers, input_image)
-    else:
-        layer_weights, layers, layer_counters = extract_layers_and_weights(
-            model, layers_to_check, sequential_order
-        )
-        feature_maps, layer_names = extract_feature_maps(
-            [layer for layer_list in layers.values() for layer in layer_list],
-            input_image,
-        )
+    activations = {}
+    # Register hooks
+
+    def get_activation(name):
+        def hook(model, input, output):
+            activations[name] = output.detach()
+        return hook
+
+    # Register hooks on the layers of interest
+    for name, module in model.named_modules():
+        if isinstance(module, tuple(layers_to_check)):
+            module.register_forward_hook(get_activation(name))
+
+    # Pass the input through the model
+    _ = model(input_image)
+
+    # Now activations dictionary has the outputs
+    feature_maps = list(activations.values())
+    layer_names = list(activations.keys())
 
     processed_feature_maps = process_feature_maps(feature_maps)
     if print_image:
         plot_feature_maps(processed_feature_maps, layer_names)
     images = []
-    for feature_map in processed_feature_maps:
+    for feature_map, layer_name in zip(processed_feature_maps, layer_names):
         if feature_map.ndim == 2:  # Grayscale feature map
             colored_feature_map = apply_colormap(feature_map)
-            images.append(colored_feature_map)
-        elif feature_map.ndim == 3 and feature_map.shape[0] == 3:  # RGB feature map
-            images.append(np.transpose(feature_map, (1, 2, 0)))
-
-        return images
+            images.append((colored_feature_map, layer_name))
+        # RGB feature map
+        elif feature_map.ndim == 3 and feature_map.shape[0] == 3:
+            images.append((np.transpose(feature_map, (1, 2, 0)), layer_name))
+    return images

PyTorchLayerViz/main.py

@@ -6,6 +6,8 @@
 import matplotlib.pyplot as plt
 
 ##########################################################
+
+
 def extract_layers_and_weights(model, layers_to_check, sequential_order=True):
     # Initialize dictionaries and lists to store layers and their weights
     layer_weights = {}
@@ -126,53 +128,61 @@ def plot_feature_maps(processed_feature_maps, layer_names):
             ax.set_title(layer_names[i].split("(")[0], fontsize=30)
             plot_index += 1
         else:
-            print(f"Skipping feature map at index {i} with shape: {feature_map.shape}")
+            print(f"Skipping feature map at index {
+                  i} with shape: {feature_map.shape}")
     plt.show()
 
 
 ##########################################################
 
 
 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, 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
+    # Load and preprocess the image
+    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(
-            model, layers_to_check, sequential_order
-        )
-        feature_maps, layer_names = extract_feature_maps(ordered_layers, input_image)
-    else:
-        layer_weights, layers, layer_counters = extract_layers_and_weights(
-            model, layers_to_check, sequential_order
-        )
-        feature_maps, layer_names = extract_feature_maps(
-            [layer for layer_list in layers.values() for layer in layer_list],
-            input_image,
-        )
+    activations = {}
+    # Register hooks
+
+    def get_activation(name):
+        def hook(model, input, output):
+            activations[name] = output.detach()
+        return hook
+
+    # Register hooks on the layers of interest
+    for name, module in model.named_modules():
+        if isinstance(module, tuple(layers_to_check)):
+            module.register_forward_hook(get_activation(name))
+
+    # Pass the input through the model
+    _ = model(input_image)
+
+    # Now activations dictionary has the outputs
+    feature_maps = list(activations.values())
+    layer_names = list(activations.keys())
 
     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:
+    for feature_map, layer_name in zip(processed_feature_maps, layer_names):
         if feature_map.ndim == 2:  # Grayscale feature map
             colored_feature_map = apply_colormap(feature_map)
-            images.append(colored_feature_map)
-        elif feature_map.ndim == 3 and feature_map.shape[0] == 3:  # RGB feature map
-            images.append(np.transpose(feature_map, (1, 2, 0)))
-
-        return images
+            images.append((colored_feature_map, layer_name))
+        # RGB feature map
+        elif feature_map.ndim == 3 and feature_map.shape[0] == 3:
+            images.append((np.transpose(feature_map, (1, 2, 0)), layer_name))
+    return images

build/lib/PyTorchLayerViz/main.py

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