enczoo 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
enczoo-0.1.0/LICENSE ADDED
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Michael J. Lee
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
enczoo-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,74 @@
1
+ Metadata-Version: 2.3
2
+ Name: enczoo
3
+ Version: 0.1.0
4
+ Summary: Map images (as `PIL.Images`) to intermediate representations (as `np.ndarray`) from off-the-shelf vision models.
5
+ Author: himjl
6
+ Author-email: himjl <mil@mit.edu>
7
+ License: MIT License
8
+
9
+ Copyright (c) 2026 Michael J. Lee
10
+
11
+ Permission is hereby granted, free of charge, to any person obtaining a copy
12
+ of this software and associated documentation files (the "Software"), to deal
13
+ in the Software without restriction, including without limitation the rights
14
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
15
+ copies of the Software, and to permit persons to whom the Software is
16
+ furnished to do so, subject to the following conditions:
17
+
18
+ The above copyright notice and this permission notice shall be included in all
19
+ copies or substantial portions of the Software.
20
+
21
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
22
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
23
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
24
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
25
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
26
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
27
+ SOFTWARE.
28
+ Requires-Dist: h5py>=3.15.1
29
+ Requires-Dist: pillow>=12.1.0
30
+ Requires-Dist: pydantic>=2.12.5
31
+ Requires-Dist: pytest>=9.0.2
32
+ Requires-Dist: torch>=2.9.1
33
+ Requires-Dist: torchvision>=0.24.1
34
+ Requires-Dist: tqdm>=4.67.1
35
+ Requires-Python: >=3.12
36
+ Description-Content-Type: text/markdown
37
+
38
+ # `enczoo`: a zoo of encoding models for images
39
+
40
+ [![CI](https://github.com/himjl/enczoo/actions/workflows/ci.yml/badge.svg)](https://github.com/himjl/enczoo/actions/workflows/ci.yml)
41
+
42
+ `enczoo` is a Python library with a single goal: to map images (as `PIL.Images`) to intermediate representations (as `np.ndarray`) from off-the-shelf vision models, such as AlexNet and ResNet50.
43
+
44
+ This library is meant for those who just need to compute off-the-shelf image features once for their project (and perhaps cache them elsewhere).
45
+
46
+ ### Installation
47
+
48
+ `enczoo` requires Python 3.12 or above, and may be installed using [uv](https://docs.astral.sh/uv/) with the following command:
49
+
50
+ >`uv add enczoo`
51
+
52
+ ### Usage
53
+
54
+ ```python
55
+ import enczoo
56
+ import PIL.Image
57
+ image = PIL.Image.open('my-image.png')
58
+
59
+ model = enczoo.ResNet50(layer_name='avgpool') # try layer4, layer3, ...
60
+ features = model.compute_features(images=[image]) # np.ndarray
61
+
62
+ # Want another layer? Check out: print(enczoo.ResNet50.layer_names)
63
+ ```
64
+
65
+ ### Things `enczoo` handles
66
+ `enczoo` aims to "just work" by solving several tiny problems which collectively make computing image features a bit annoying. `enczoo` handles:
67
+
68
+ * performing model-specific image normalization ("_was it -1 to 1, 0 to 1, 0-255...? ImageNet channel normalization...?_"),
69
+ * correctly encoding images ("_my image was in mode L, not RGB!_")
70
+ * turning off any batch normalization ("_was the model in training mode...?_")
71
+ * extracting intermediate layers by name ("_how do I do that forward hook thing again...?_")
72
+ * turning off autograd, and returning tensors as `np.ndarray` (no more `.cpu().numpy()`)
73
+ * image cropping to fit input tensor shape (default: center cropping. no black bars!)
74
+ * and more!
enczoo-0.1.0/README.md ADDED
@@ -0,0 +1,37 @@
1
+ # `enczoo`: a zoo of encoding models for images
2
+
3
+ [![CI](https://github.com/himjl/enczoo/actions/workflows/ci.yml/badge.svg)](https://github.com/himjl/enczoo/actions/workflows/ci.yml)
4
+
5
+ `enczoo` is a Python library with a single goal: to map images (as `PIL.Images`) to intermediate representations (as `np.ndarray`) from off-the-shelf vision models, such as AlexNet and ResNet50.
6
+
7
+ This library is meant for those who just need to compute off-the-shelf image features once for their project (and perhaps cache them elsewhere).
8
+
9
+ ### Installation
10
+
11
+ `enczoo` requires Python 3.12 or above, and may be installed using [uv](https://docs.astral.sh/uv/) with the following command:
12
+
13
+ >`uv add enczoo`
14
+
15
+ ### Usage
16
+
17
+ ```python
18
+ import enczoo
19
+ import PIL.Image
20
+ image = PIL.Image.open('my-image.png')
21
+
22
+ model = enczoo.ResNet50(layer_name='avgpool') # try layer4, layer3, ...
23
+ features = model.compute_features(images=[image]) # np.ndarray
24
+
25
+ # Want another layer? Check out: print(enczoo.ResNet50.layer_names)
26
+ ```
27
+
28
+ ### Things `enczoo` handles
29
+ `enczoo` aims to "just work" by solving several tiny problems which collectively make computing image features a bit annoying. `enczoo` handles:
30
+
31
+ * performing model-specific image normalization ("_was it -1 to 1, 0 to 1, 0-255...? ImageNet channel normalization...?_"),
32
+ * correctly encoding images ("_my image was in mode L, not RGB!_")
33
+ * turning off any batch normalization ("_was the model in training mode...?_")
34
+ * extracting intermediate layers by name ("_how do I do that forward hook thing again...?_")
35
+ * turning off autograd, and returning tensors as `np.ndarray` (no more `.cpu().numpy()`)
36
+ * image cropping to fit input tensor shape (default: center cropping. no black bars!)
37
+ * and more!
@@ -0,0 +1,32 @@
1
+ [project]
2
+ name = "enczoo"
3
+ version = "0.1.0"
4
+ description = "Map images (as `PIL.Images`) to intermediate representations (as `np.ndarray`) from off-the-shelf vision models."
5
+ readme = "README.md"
6
+ authors = [
7
+ { name = "himjl", email = "mil@mit.edu" }
8
+ ]
9
+ license = { file = "LICENSE" }
10
+ requires-python = ">=3.12"
11
+ dependencies = [
12
+ "h5py>=3.15.1",
13
+ "pillow>=12.1.0",
14
+ "pydantic>=2.12.5",
15
+ "pytest>=9.0.2",
16
+ "torch>=2.9.1",
17
+ "torchvision>=0.24.1",
18
+ "tqdm>=4.67.1",
19
+ ]
20
+
21
+ [build-system]
22
+ requires = ["uv_build>=0.8.12,<0.9.0"]
23
+ build-backend = "uv_build"
24
+
25
+ [dependency-groups]
26
+ dev = [
27
+ "black>=25.12.0",
28
+ "matplotlib>=3.10.8",
29
+ "pytest>=9.0.2",
30
+ "ruff>=0.14.10",
31
+ "ty>=0.0.9",
32
+ ]
@@ -0,0 +1,6 @@
1
+ __all__ = ["ImageEncoding", "Pixels", "ResNet50", "AlexNet"]
2
+
3
+
4
+ from enczoo.base import ImageEncoding
5
+ from enczoo.neural_networks.pretrained import AlexNet, ResNet50
6
+ from enczoo.classic.pixels import Pixels
@@ -0,0 +1,172 @@
1
+ from abc import ABC, abstractmethod
2
+ from typing import List, Tuple
3
+
4
+ import PIL.Image
5
+ import numpy as np
6
+ import torch
7
+
8
+ import enczoo.utils as utils
9
+
10
+
11
+ # %%
12
+ class ImageEncoding(
13
+ torch.nn.Module,
14
+ ABC,
15
+ ):
16
+ """Map PIL images to batched float tensors.
17
+
18
+ This module maps B-length lists of PIL.Image.Image to float tensors shaped
19
+ [B, *]. Parameters do not aggregate gradients by default.
20
+ """
21
+
22
+ def __init__(
23
+ self,
24
+ trainable: bool = False,
25
+ ):
26
+ """Initialize the encoder.
27
+
28
+ Args:
29
+ trainable: If True, the module is put in train mode.
30
+ """
31
+ super().__init__()
32
+ self.trainable = trainable
33
+ self._module_hash = None
34
+ self._tensor_bucket = None
35
+ self._output_shape = None
36
+
37
+ # Set the module's mode
38
+ self.train(mode=trainable)
39
+
40
+ @property
41
+ def device(self) -> torch.device:
42
+ """Infer the device from the first parameter or buffer.
43
+
44
+ Returns:
45
+ The device of the first parameter or buffer, or CPU if none exist.
46
+ """
47
+ try:
48
+ return next(self.parameters()).device
49
+ except StopIteration:
50
+ try:
51
+ return next(self.buffers()).device
52
+ except StopIteration:
53
+ return torch.device("cpu")
54
+
55
+ @property
56
+ def output_shape(self) -> Tuple[int, ...]:
57
+ """Return the output feature shape (excluding batch dimension)."""
58
+ if self._output_shape is None:
59
+ test_image = PIL.Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))
60
+ test_result = self.compute_features(images=[test_image], flatten=False)
61
+ if not isinstance(test_result, np.ndarray):
62
+ raise ValueError(
63
+ f"Expected a np.ndarray from self.forward, but got {type(test_result)}"
64
+ )
65
+ if not test_result.shape[0] == 1:
66
+ raise ValueError(
67
+ f"Expected a batch size of 1, but got {test_result.shape}"
68
+ )
69
+ if len(test_result.shape) == 1:
70
+ output_shape = tuple()
71
+ else:
72
+ output_shape = test_result.shape[1:]
73
+
74
+ self._output_shape = tuple(output_shape)
75
+
76
+ return self._output_shape
77
+
78
+ @property
79
+ def module_hash(self) -> str:
80
+ """Return a stable hash for the module's parameters and structure."""
81
+ if self.trainable:
82
+ raise ValueError("Cannot hash a trainable model.")
83
+
84
+ # Turn off gradients for all parameters
85
+ for param in self.parameters():
86
+ param.requires_grad = False
87
+
88
+ # Hash self if unhashed
89
+ if self._module_hash is None:
90
+ self._module_hash = utils.hash_torch_module(module=self)
91
+ return self._module_hash
92
+
93
+ def compute_features(
94
+ self,
95
+ images: List[PIL.Image.Image],
96
+ flatten: bool = False,
97
+ seed: int | None = None,
98
+ ) -> np.ndarray:
99
+ """Compute features and return them as a NumPy array.
100
+
101
+ This is an alias for __call__ to help IDEs that do not recognize the
102
+ torch.nn.Module __call__ signature.
103
+
104
+ Args:
105
+ images: A B-length list of PIL.Image.Image.
106
+ flatten: If True, flatten the output to [B, d].
107
+ seed: Optional RNG seed for deterministic results.
108
+
109
+ Returns:
110
+ A NumPy array of shape [B, *], or [B, d] if flatten=True.
111
+
112
+ Raises:
113
+ ValueError: If the input images are invalid.
114
+ """
115
+
116
+ with torch.random.fork_rng():
117
+ if seed is not None:
118
+ torch.manual_seed(seed)
119
+ if torch.cuda.is_available():
120
+ torch.cuda.manual_seed_all(seed)
121
+ with torch.no_grad():
122
+ torch_features: torch.Tensor = self(images=images, flatten=flatten)
123
+ numpy_features = torch_features.detach().cpu().numpy()
124
+ return numpy_features
125
+
126
+ def forward(
127
+ self,
128
+ images: List[PIL.Image.Image],
129
+ flatten: bool = False,
130
+ ) -> torch.Tensor:
131
+ """Compute features for a batch of images.
132
+
133
+ Args:
134
+ images: A B-length list of PIL.Image.Image.
135
+ flatten: If True, flatten the output to [B, d].
136
+
137
+ Returns:
138
+ A torch.Tensor of shape [B, *].
139
+
140
+ Raises:
141
+ ValueError: If the images list is empty or not image objects.
142
+ """
143
+ if not isinstance(images, list):
144
+ raise ValueError(
145
+ f"Expected a list of PIL.Image.Images, but got {type(images)}"
146
+ )
147
+ if len(images) == 0:
148
+ raise ValueError("Expected a non-empty list of PIL.Image.Images.")
149
+ if not isinstance(images[0], PIL.Image.Image):
150
+ raise ValueError(
151
+ f"Expected a list of PIL.Image.Images, but element 0 is a {type(images[0])}"
152
+ )
153
+
154
+ # Call the subclass implementation
155
+ feats = self._images_to_features(images=images)
156
+ if flatten:
157
+ # Flatten the features
158
+ feats = feats.reshape(feats.shape[0], -1)
159
+
160
+ return feats
161
+
162
+ @abstractmethod
163
+ def _images_to_features(self, images: List[PIL.Image.Image]) -> torch.Tensor:
164
+ """Convert images to features.
165
+
166
+ Args:
167
+ images: A list of PIL.Image.Image.
168
+
169
+ Returns:
170
+ A torch.Tensor of shape [B, *].
171
+ """
172
+ raise NotImplementedError
File without changes
@@ -0,0 +1,89 @@
1
+ from typing import List
2
+
3
+ import PIL.Image
4
+ import torch
5
+ import torchvision.transforms.v2 as v2
6
+
7
+ from enczoo.base import ImageEncoding
8
+ from enczoo.transforms.random_projection import RandomProjection
9
+
10
+
11
+ # %%
12
+ class Pixels(ImageEncoding):
13
+ """Encode images by their resized center-crop pixels."""
14
+
15
+ def __init__(
16
+ self,
17
+ size: int = 16,
18
+ random_projection_dim: int | None = None,
19
+ random_projection_seed: int | None = None,
20
+ ):
21
+ """Initialize the pixel encoder.
22
+
23
+ Args:
24
+ size: Output side length in pixels.
25
+ random_projection_dim: Optional output dimension for projection.
26
+ random_projection_seed: Seed for projection weights.
27
+ """
28
+ super().__init__()
29
+
30
+ # Register size tensor as buffer
31
+ self.register_buffer(
32
+ "size", torch.tensor(size, dtype=torch.int16, requires_grad=False)
33
+ )
34
+
35
+ # Transform
36
+ self.transforms = v2.Compose(
37
+ [
38
+ v2.ToImage(),
39
+ v2.RGB(),
40
+ v2.ToDtype(torch.uint8, scale=True),
41
+ v2.Resize(size=None, max_size=size, antialias=False),
42
+ v2.CenterCrop(size=size),
43
+ v2.ToDtype(torch.float32, scale=True),
44
+ ]
45
+ )
46
+
47
+ # Random projection
48
+ if random_projection_dim is not None:
49
+ if random_projection_dim > (size * size * 3):
50
+ raise ValueError(
51
+ f"random_projection_dim must be less than or equal to size * size * 3={size * size * 3}!"
52
+ )
53
+
54
+ if random_projection_seed is None:
55
+ raise ValueError(
56
+ "random_projection_seed must be provided if random_projection_dim is not None!"
57
+ )
58
+ self.random_projection = RandomProjection(
59
+ seed=random_projection_seed,
60
+ in_features=int(size * size * 3),
61
+ out_features=random_projection_dim,
62
+ )
63
+ else:
64
+ self.random_projection = None
65
+
66
+ def _images_to_features(self, images: List[PIL.Image.Image]) -> torch.Tensor:
67
+ """Convert images to pixel features.
68
+
69
+ Args:
70
+ images: A list of PIL.Image.Image.
71
+
72
+ Returns:
73
+ A torch.Tensor of shape [B, size, size, 3] or projected features.
74
+ """
75
+ # Apply the transformations to each image
76
+ transformed_images = [self.transforms(image.convert("RGB")) for image in images]
77
+
78
+ # Stack the transformed images into a single tensor
79
+ images_tensor = torch.stack(transformed_images)
80
+
81
+ # If random projection is enabled, apply it
82
+ if self.random_projection is not None:
83
+ return self.random_projection(
84
+ images_tensor.reshape(images_tensor.shape[0], -1)
85
+ )
86
+ else:
87
+ # Rearrange from BCHW to BHWC order
88
+ images_tensor = images_tensor.permute(0, 2, 3, 1)
89
+ return images_tensor
File without changes
@@ -0,0 +1,184 @@
1
+ from abc import ABC
2
+ from typing import Dict, List, Tuple
3
+
4
+ import PIL.Image
5
+ import numpy as np
6
+ import torch
7
+ import torchvision
8
+
9
+ from enczoo.base import ImageEncoding
10
+ from enczoo.transforms.random_projection import RandomProjection
11
+
12
+
13
+ class ImageNeuralNetwork(ImageEncoding, ABC):
14
+ """Image encoding backed by a torch neural network."""
15
+
16
+ def __init__(
17
+ self,
18
+ image_loader: torch.nn.Module | torchvision.transforms.Compose,
19
+ model: torch.nn.Module,
20
+ layer_name: str,
21
+ random_projection_dim: int | None,
22
+ random_projection_seed: int | None,
23
+ ):
24
+ """Initialize the neural network encoder.
25
+
26
+ Args:
27
+ image_loader: Module that converts PIL images to model inputs.
28
+ model: Torch model used to compute activations.
29
+ layer_name: Name of the layer whose activations are returned.
30
+ random_projection_dim: Optional output dimension for projection.
31
+ random_projection_seed: Seed for projection weights.
32
+
33
+ Raises:
34
+ ValueError: If the layer name is not found.
35
+ """
36
+ super().__init__()
37
+
38
+ # Ensure modules will be registered in evaluation mode
39
+ self.train(mode=False)
40
+
41
+ # Register buffers to ensure the model's hash is distinctive for each layer
42
+ self.register_buffer(
43
+ "layer_name", torch.tensor([ord(c) for c in layer_name], dtype=torch.int16)
44
+ )
45
+ self._layer_name = layer_name # Needed for the forward pass
46
+
47
+ def register_hook(
48
+ module: torch.nn.Module,
49
+ root_name: str,
50
+ activations_dict: Dict[str, torch.Tensor],
51
+ ) -> List[str]:
52
+ """Recursively register forward hooks on leaf modules.
53
+
54
+ Args:
55
+ module: Module whose children are walked.
56
+ root_name: Prefix for module names.
57
+ activations_dict: Dict populated with layer activations.
58
+
59
+ Returns:
60
+ A list of layer names in discovery order.
61
+ """
62
+ module_names = []
63
+
64
+ nchildren = 0
65
+ for module_name, submodule in module.named_children():
66
+ if module_name != "":
67
+ next_root_name = (
68
+ root_name + "." + module_name
69
+ if root_name != ""
70
+ else module_name
71
+ )
72
+ else:
73
+ raise ValueError("Empty module name found in model!")
74
+ # Ensure module is in evaluation mode
75
+ submodule.train(mode=False)
76
+ # Recursive call:
77
+ submodule_names = register_hook(
78
+ submodule,
79
+ root_name=next_root_name,
80
+ activations_dict=activations_dict,
81
+ )
82
+
83
+ # Update the number of children
84
+ nchildren += 1
85
+ module_names.extend(submodule_names)
86
+
87
+ # Base case:
88
+ if nchildren == 0:
89
+ layer_name = root_name
90
+
91
+ if layer_name in activations_dict:
92
+ # Don't think this should ever happen:
93
+ raise Exception(
94
+ f"Layer name {layer_name} already exists in hidden activations! Existing keys: {self._hidden_activations.keys()}"
95
+ )
96
+
97
+ def hook_function(module: torch.nn.Module, args, output):
98
+ # print(f'Hook called on {layer_name}')
99
+ activations_dict[layer_name] = output
100
+
101
+ # Attach forward hook
102
+ module.register_forward_hook(hook_function)
103
+
104
+ # Base case: no children
105
+ module_names.append(layer_name)
106
+ return module_names
107
+
108
+ # Register forward hooks that will populate this dictionary with hidden activations on the forward pass:
109
+ self._hidden_activations: Dict[str, torch.Tensor] = {}
110
+ self._layer_names = register_hook(
111
+ model, root_name="", activations_dict=self._hidden_activations
112
+ )
113
+
114
+ if layer_name not in self._layer_names:
115
+ raise ValueError(
116
+ f"Layer name {layer_name} not found in model.\nAvailable layer names: {self._layer_names}"
117
+ )
118
+
119
+ # Populate the sizes of the layers with a forward pass
120
+ with torch.no_grad():
121
+ test_image = PIL.Image.fromarray(np.zeros((224, 224, 3), dtype=np.uint8))
122
+ test_image = image_loader(test_image)
123
+ model(test_image.unsqueeze(0))
124
+
125
+ self._layer_to_shape = {
126
+ layer: tuple(self._hidden_activations[layer].shape[1:])
127
+ for layer in self._hidden_activations
128
+ }
129
+
130
+ # Register modules
131
+ self.image_loader = image_loader
132
+ self.model = model
133
+
134
+ if random_projection_dim is not None:
135
+ if random_projection_seed is None:
136
+ raise ValueError(
137
+ "random_projection_seed must be provided if random_projection_dim is not None!"
138
+ )
139
+ self.random_projection = RandomProjection(
140
+ seed=random_projection_seed,
141
+ in_features=int(np.prod(self._layer_to_shape[layer_name])),
142
+ out_features=random_projection_dim,
143
+ )
144
+ else:
145
+ self.random_projection = None
146
+
147
+ def _images_to_features(self, images: List[PIL.Image.Image]) -> torch.Tensor:
148
+ """Convert images to network activations.
149
+
150
+ Args:
151
+ images: A list of PIL.Image.Image.
152
+
153
+ Returns:
154
+ A torch.Tensor of shape [B, *].
155
+ """
156
+
157
+ # Preprocess the images
158
+ preprocessed_images = torch.stack(
159
+ [self.image_loader(image) for image in images], dim=0
160
+ )
161
+
162
+ # Transfer to the correct device
163
+ preprocessed_images = preprocessed_images.to(self.device)
164
+
165
+ # Run the forward pass
166
+ self.model(preprocessed_images)
167
+
168
+ # Retrieve the activations for the given layer
169
+ f = self._hidden_activations[self._layer_name]
170
+
171
+ # Perform random projection if requested
172
+ if self.random_projection is not None:
173
+ # Flatten the features
174
+ f = f.reshape(f.shape[0], -1)
175
+
176
+ # Run the random projection forward
177
+ f = self.random_projection(f)
178
+
179
+ return f
180
+
181
+ @property
182
+ def layer_name_to_shape(self) -> Dict[str, Tuple[int, ...]]:
183
+ """Return a mapping of layer names to activation shapes."""
184
+ return self._layer_to_shape
@@ -0,0 +1,163 @@
1
+ from abc import ABC, abstractmethod
2
+ from typing import List, Tuple
3
+
4
+ import PIL.Image
5
+ import torch
6
+ import torchvision.models
7
+ import torchvision.transforms.functional as F
8
+
9
+ from enczoo.neural_networks.base import ImageNeuralNetwork
10
+
11
+
12
+ class StandardImageLoader(torch.nn.Module):
13
+ """Load and normalize images for standard torchvision models.
14
+
15
+ This loader resizes directly to 224 before center-crop, preserving the
16
+ largest square sub-image. It also converts inputs to RGB.
17
+
18
+ Example:
19
+ - Original loader: resize to 256, then center-crop 224.
20
+ - This loader: resize to 224, then center-crop 224.
21
+ """
22
+
23
+ def forward(self, img: PIL.Image.Image) -> torch.Tensor:
24
+ """Convert a PIL image into a normalized tensor.
25
+
26
+ Args:
27
+ img: Input image.
28
+
29
+ Returns:
30
+ Normalized image tensor suitable for torchvision models.
31
+ """
32
+ img = img.convert("RGB")
33
+
34
+ img_tensor = F.pil_to_tensor(pic=img)
35
+ img_tensor = F.resize(
36
+ img=img_tensor, size=[224], interpolation=F.InterpolationMode.BILINEAR
37
+ )
38
+ img_tensor = F.center_crop(img=img_tensor, output_size=[224])
39
+ img_tensor = F.convert_image_dtype(image=img_tensor, dtype=torch.float)
40
+ img_tensor = F.normalize(
41
+ tensor=img_tensor,
42
+ mean=[0.485, 0.456, 0.406],
43
+ std=[0.229, 0.224, 0.225],
44
+ )
45
+ return img_tensor
46
+
47
+
48
+ class _PretrainedNN(ImageNeuralNetwork, ABC):
49
+ """Base class for pretrained torchvision encoders."""
50
+
51
+ layer_names: List[str]
52
+
53
+ def __init__(
54
+ self,
55
+ layer_name: str,
56
+ random_projection_dim: int | None = None,
57
+ random_projection_seed: int | None = None,
58
+ ):
59
+ """Initialize a pretrained encoder.
60
+
61
+ Args:
62
+ layer_name: Name of the layer whose activations are returned.
63
+ random_projection_dim: Optional output dimension for projection.
64
+ random_projection_seed: Seed for projection weights.
65
+
66
+ Raises:
67
+ ValueError: If the layer name is invalid.
68
+ """
69
+ if layer_name not in self.layer_names:
70
+ raise ValueError(
71
+ f"Unknown layer_name: {layer_name}. Available:\n{self.layer_names}"
72
+ )
73
+
74
+ image_loader, model = self._load_modules()
75
+
76
+ # Ensure modules are in evaluation mode by default
77
+ image_loader.train(mode=False)
78
+ model.train(mode=False)
79
+
80
+ super().__init__(
81
+ image_loader=image_loader,
82
+ model=model,
83
+ layer_name=layer_name,
84
+ random_projection_dim=random_projection_dim,
85
+ random_projection_seed=random_projection_seed,
86
+ )
87
+
88
+ @abstractmethod
89
+ def _load_modules(self) -> Tuple[torch.nn.Module, torch.nn.Module]:
90
+ """Load the image loader and model for this network.
91
+
92
+ Returns:
93
+ A tuple of (image_loader, model).
94
+ """
95
+ raise NotImplementedError
96
+
97
+
98
+ class AlexNet(_PretrainedNN):
99
+ """AlexNet encoder with named layer outputs."""
100
+
101
+ # A subset of all layers (each separated by one nonlinearity):
102
+ layer_names = [
103
+ "features.1",
104
+ "features.4",
105
+ "features.7",
106
+ "features.9",
107
+ "features.11",
108
+ "classifier.2",
109
+ "classifier.5",
110
+ "classifier.6",
111
+ ]
112
+
113
+ def _load_modules(self):
114
+ """Load the AlexNet image loader and model."""
115
+ image_loader = StandardImageLoader()
116
+ model = torchvision.models.alexnet(
117
+ weights=torchvision.models.AlexNet_Weights.IMAGENET1K_V1
118
+ )
119
+ return image_loader, model
120
+
121
+
122
+ class ResNet50(_PretrainedNN):
123
+ """ResNet-50 encoder with named layer outputs."""
124
+
125
+ # A subset of layers (each separated by one nonlinearity, except layer4.2.relu, avgpool, and fc, which are connected by a linear layer):
126
+ layer_names = [
127
+ "relu",
128
+ "layer1.0.relu",
129
+ "layer1.1.relu",
130
+ "layer1.2.relu",
131
+ "layer2.0.relu",
132
+ "layer2.1.relu",
133
+ "layer2.2.relu",
134
+ "layer2.3.relu",
135
+ "layer3.0.relu",
136
+ "layer3.1.relu",
137
+ "layer3.2.relu",
138
+ "layer3.3.relu",
139
+ "layer3.4.relu",
140
+ "layer3.5.relu",
141
+ "layer4.0.relu",
142
+ "layer4.1.relu",
143
+ "layer4.2.relu",
144
+ "avgpool",
145
+ "fc",
146
+ ]
147
+
148
+ def _load_modules(self):
149
+ """Load the ResNet-50 image loader and model."""
150
+ image_loader = StandardImageLoader()
151
+ model = torchvision.models.resnet50(
152
+ weights=torchvision.models.ResNet50_Weights.IMAGENET1K_V1
153
+ )
154
+ return image_loader, model
155
+
156
+
157
+ if __name__ == "__main__":
158
+ resnet50 = ResNet50(
159
+ layer_name="avgpool", random_projection_dim=None, random_projection_seed=0
160
+ )
161
+ print(resnet50.training)
162
+ print(resnet50.model.training)
163
+ print(getattr(resnet50.image_loader, "training", None))
@@ -0,0 +1,54 @@
1
+ from typing import List
2
+
3
+ import PIL.Image
4
+ import numpy as np
5
+ import pytest
6
+ import torch
7
+
8
+ from enczoo import AlexNet
9
+
10
+
11
+ @pytest.fixture
12
+ def alexnet():
13
+ return AlexNet(
14
+ layer_name=AlexNet.layer_names[-1],
15
+ random_projection_dim=None,
16
+ random_projection_seed=0,
17
+ )
18
+
19
+
20
+ @pytest.fixture
21
+ def images() -> List[PIL.Image.Image]:
22
+ np.random.seed(0)
23
+ img_dat1 = PIL.Image.fromarray(
24
+ np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
25
+ )
26
+ img_dat2 = PIL.Image.fromarray(
27
+ np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
28
+ )
29
+ return [img_dat1, img_dat2]
30
+
31
+
32
+ @pytest.fixture
33
+ def image_batch1(images):
34
+ return [images[0]]
35
+
36
+
37
+ @pytest.fixture
38
+ def image_batch2(images):
39
+ return images
40
+
41
+
42
+ def test_default_is_not_training(alexnet):
43
+ assert not alexnet.training
44
+
45
+
46
+ def test_alexnet_forward(alexnet, image_batch1, image_batch2):
47
+ # Test output is deterministic
48
+ result = alexnet(images=image_batch1)
49
+ result2 = alexnet(images=image_batch1)
50
+ assert torch.allclose(result, result2)
51
+
52
+ # Test output does not depend on batch size
53
+ result_bigger_batch = alexnet(images=image_batch2)
54
+ assert torch.allclose(result[0], result_bigger_batch[0], rtol=1e-3)
File without changes
File without changes
@@ -0,0 +1,73 @@
1
+ import math
2
+ from typing import Tuple, cast
3
+
4
+ import torch
5
+ import torch.nn as nn
6
+ import torch.nn.functional as F
7
+
8
+
9
+ class RandomProjection(nn.Module):
10
+ """Apply a fixed random projection to an input tensor."""
11
+
12
+ def __init__(
13
+ self,
14
+ in_features: int,
15
+ out_features: int,
16
+ seed: int,
17
+ ):
18
+ """Initialize the random projection.
19
+
20
+ Args:
21
+ in_features: Input feature dimension.
22
+ out_features: Output feature dimension.
23
+ seed: Seed for the random projection weights.
24
+ """
25
+
26
+ super().__init__()
27
+ self.train(mode=False)
28
+
29
+ # Register inputs as buffers; these will constitute the module's hash.
30
+ self.register_buffer(
31
+ "seed", torch.tensor(seed, dtype=torch.int64, requires_grad=False)
32
+ )
33
+ self.register_buffer(
34
+ "in_features",
35
+ torch.tensor(in_features, dtype=torch.int64, requires_grad=False),
36
+ )
37
+ self.register_buffer(
38
+ "out_features",
39
+ torch.tensor(out_features, dtype=torch.int64, requires_grad=False),
40
+ )
41
+
42
+ with torch.random.fork_rng():
43
+ # Set the weights from a standard normal distribution:
44
+ torch.manual_seed(seed)
45
+ weights = torch.randn(
46
+ size=(out_features, in_features),
47
+ dtype=torch.float32,
48
+ requires_grad=False,
49
+ ) / math.sqrt(in_features * out_features)
50
+
51
+ self.register_buffer("projection_weights", weights)
52
+
53
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
54
+ """Project input features with a fixed random matrix.
55
+
56
+ Args:
57
+ x: Input tensor of shape [B, in_features].
58
+
59
+ Returns:
60
+ Projected tensor of shape [B, out_features].
61
+ """
62
+ weights = cast(torch.Tensor, self.projection_weights)
63
+ return F.linear(x, weights)
64
+
65
+ def __repr__(self):
66
+ """Return a concise representation for debugging."""
67
+ return f"RandomProjection(in_features={self.in_features}, out_features={self.out_features}, seed={self.seed})"
68
+
69
+ @property
70
+ def output_shape(self) -> Tuple[int]:
71
+ """Return the output feature shape (excluding batch dimension)."""
72
+ weights = cast(torch.Tensor, self.projection_weights)
73
+ return (weights.shape[0],)
@@ -0,0 +1,60 @@
1
+ import hashlib
2
+ import itertools
3
+ from typing import Iterable, Iterator, List, TypeVar
4
+
5
+ import torch
6
+
7
+ T = TypeVar("T")
8
+
9
+
10
+ def iterate_batches(iterable: Iterable[T], batch_size: int) -> Iterator[List[T]]:
11
+ """Yield items from an iterable in fixed-size batches.
12
+
13
+ Args:
14
+ iterable: Source iterable to batch.
15
+ batch_size: Number of items per batch.
16
+
17
+ Yields:
18
+ Lists of up to batch_size elements.
19
+
20
+ Raises:
21
+ ValueError: If batch_size is less than 1.
22
+ """
23
+ if batch_size < 1:
24
+ raise ValueError("Batch size must be at least 1!")
25
+
26
+ iterator = iter(iterable)
27
+ while True:
28
+ batch = list(itertools.islice(iterator, batch_size))
29
+ if len(batch) == 0:
30
+ return
31
+ yield batch
32
+
33
+
34
+ def hash_torch_module(module: torch.nn.Module) -> str:
35
+ """Return a hash for a torch.nn.Module.
36
+
37
+ The hash depends on the module's state_dict and string representation.
38
+
39
+ Args:
40
+ module: Module to hash.
41
+
42
+ Returns:
43
+ Hex-encoded SHA-256 hash.
44
+ """
45
+
46
+ sha256_hash = hashlib.sha256()
47
+
48
+ # Hash the model's state_dict:
49
+ state_dict = module.state_dict()
50
+ for key in sorted(state_dict.keys()):
51
+ tensor_value = state_dict[key].detach().cpu().numpy()
52
+ sha256_hash.update(tensor_value.tobytes())
53
+
54
+ # Hash the module's string representation:
55
+ module_string = str(module)
56
+ module_string = module_string.encode("utf-8")
57
+ sha256_hash.update(module_string)
58
+
59
+ # Return the combined hash:
60
+ return sha256_hash.hexdigest()