Skip to content

Latest commit

 

History

History
241 lines (163 loc) · 13.6 KB

README.md

File metadata and controls

241 lines (163 loc) · 13.6 KB
Raw

stylemod

PyPI - Version PyPI - License

Modular neural style transfer (NST) library designed to make it easy to integrate and customize different deep learning models for artistic style transfer.

Table of Contents

Key Features

  • Plug-and-play architecture for integrating new models.
  • Support for CNN-based and Transformer-based models.
  • Easy customization of style and content loss computation.
  • Command-line interface (CLI) for easy interaction.
  • Provides out-of-the-box functionality for managing models, utilized layers/weights, normalizations, and more.

Modular Architecture

Here is a visualization of the class hierarchy for the stylemod library:

Class Hierarchy

Installation

Option 1: Install via PyPI (Recommended)

pip install stylemod

This will automatically install required dependencies. You can also view the package on PyPI.

Option 2: Install from Source (For Development)

  1. Clone the repository:

    git clone https://github.com/ooojustin/stylemod.git
cd stylemod

  2. Install dependencies: Make sure you have PyTorch and other required libraries installed:

    pip install -r requirements.txt

Install Graphviz (Optional)

If you wish to use the built-in Graphviz integration for architecture visualization, ensure Graphviz is installed:

  • Windows
    You can download Graphviz for Windows from the official website:
    Windows Download

    Alternatively, you can install it using popular package managers:

    # Using Chocolatey
choco install graphviz

# Using Scoop
scoop install graphviz

  • Unix-based Systems

    # For Linux (Debian/Ubuntu)
sudo apt-get install graphviz

# For Linux (Red Hat/CentOS)
sudo yum install graphviz

# For macOS
brew install graphviz

Note: If you try to invoke stylemod.generate_class_hierarchy() or model.visualize() without graphviz installed, stylemod will attempt to install it automatically via your package manager on Linux/MacOS.

Model Superclasses

In the stylemod library, models used for neural style transfer are designed to be modular and extensible. They inherit from two primary classes: AbstractBaseModel, which provides a blueprint for all models, and BaseModel, which extends AbstractBaseModel to provide common functionality for most neural style transfer tasks. Subclasses like CNNBaseModel and TransformerBaseModel extend BaseModel with architecture-specific logic.

AbstractBaseModel

The AbstractBaseModel is an abstract class that defines the required interface for all neural style transfer models. It does not provide any concrete implementations but instead acts as a blueprint to ensure that all models follow a consistent structure. Each model must implement methods for initialization, feature extraction, loss calculation, and visualization.

Below is a table summarizing the key abstract methods that subclasses must implement:

Abstract Method Description
initialize_module() Initializes the model architecture and loads any required weights.
get_model_module() Returns the initialized model, ensuring that it has been properly set up.
eval() Switches the model to evaluation mode, disabling training-specific operations (like dropout or batch normalization).
set_device(device) Moves the model to the specified device (CPU/GPU).
normalize_tensor(tensor) Normalizes the input tensor according to the model’s pre-defined normalization (if applicable).
denormalize_tensor(tensor) Reverts normalization applied to a tensor, returning it to its original scale and distribution.
get_features(image, layers) Extracts feature maps from the given image at specified model layers.
calc_gram_matrix(tensor) Calculates the gram matrix of a tensor, which is used to capture style information in style transfer models.
calc_content_loss(target, content_features) Computes the content loss by comparing the target image's features to the content image’s features.
calc_style_loss(target, style_features) Computes the style loss by comparing the target image's style features with those from the style image.
forward(target, content_image, style_image, content_features, style_features) Combines content and style losses into a single scalar value for optimization.
visualize() Visualizes the model’s architecture, typically outputting a Graphviz diagram.

BaseModel

The BaseModel class extends AbstractBaseModel and provides core functionality such as model initialization, normalization, feature extraction, and content/style loss computation. This class is designed to reduce repetitive code, allowing subclasses to focus on model-specific logic.

  • Initialization: The model is initialized with a callable function (model_fn) to load the architecture and optional pre-trained weights.
  • Normalization: Handles input tensor normalization and denormalization, ensuring consistent image processing.
  • Feature Extraction: Extracts feature maps from intermediate layers of the model.
  • Gram Matrix Calculation: Provides a default implementation to calculate gram matrices, used for style transfer tasks.
  • Content and Style Loss: Implements methods for calculating content and style losses based on feature maps and gram matrices.

CNNBaseModel

The CNNBaseModel extends BaseModel without overriding the content and style loss calculations, meaning it leverages the same base implementation for both loss functions. The base calc_content_loss compares content features, and the calc_style_loss compares the gram matrices of style features, making this class suitable for CNN-based neural style transfer models.

TransformerBaseModel

The TransformerBaseModel extends BaseModel to support transformer architectures that rely on attention mechanisms. This class introduces additional functionality for attention-based style transfer. When use_attention is set to True, it utilizes attention maps during style loss calculation.

  • Attention Mechanism: Requires an implementation of get_attention(), as the attention mechanism varies across different transformer architectures.
  • Style Loss: Uses both feature-based and attention-based style loss by comparing the gram matrices of feature maps and attention maps.
  • Dynamic Control: Attention-based style loss is only applied if the use_attention flag is set to True and if get_attention() is implemented in the subclass.

CLI Usage

The stylemod library provides a command-line interface (CLI) for running style transfer and visualizing model architectures.

Available Commands

1. Running Style Transfer

You can run the style transfer directly from the CLI by providing the paths to your content and style images, along with other optional parameters like the output filename, the number of steps, and the model to use.

stylemod run --content-image "content.png" --style-image "style.png" --steps 500 --model VGG19
Options:
  • -ci, --content-image: (Required) Path to the content image.
  • -si, --style-image: (Required) Path to the style image.
  • -o, --output-image: Filename for the output image. (Default: output_image.png)
  • -s, --steps: Number of optimization steps. (Default: 1000)
  • -ms, --max-size: Maximum size of input images. (Default: 400)
  • -m, --model: The model to use for style transfer. (Default: VGG19)
  • -gpu, --gpu-index: GPU index to use. (Default: 0, if available)

2. Visualizing Model Architecture

You can visualize the architecture of a specific model using the visualize command:

stylemod visualize VGG19 --output "model_visualization.png" --dpi 300
Options:
  • {model name}: The model architecture to visualize.
  • -o, --output: Optional path to save the visualization image (e.g., model_vis.png).
  • -d, --dpi: Set the DPI (dots per inch) for the rendered image. (Default: 400)

To see an example of what the output of visualizing VGG19 would look like, see visualize_vgg19.png.

3. Visualizing Class Hierarchy

You can visualize the class hierarchy of the stylemod library, which shows the relationships between different model classes.

stylemod class-hierarchy --save --show-funcs
Options:
  • -s, --save: Save the rendered class hierarchy to a file (img/class_hierarchy.png).
  • -f, --show-funcs: Show the abstract functions that should be implemented by subclasses.
  • -d, --dpi: Set the DPI (dots per inch) for the rendered image. (Default: 200)

ModelFactory

The ModelFactory class is responsible for dynamically creating instances of models used in the stylemod library. It provides a flexible and extensible mechanism for handling different model architectures and implementation without needing to hard-code their instantiations.

The ModelFactory automatically registers any model that extends AbstractBaseModel found in the stylemod.models package. Additional models can be registered manually if needed.

Key Features:

  • Dynamic Model Creation: Allows creating model instances by name or enum value, where **kwargs are forwarded to the constructor via the create() method.
  • Automatic Model Registration: Automatically scans and registers all models in the stylemod.models package that inherit from AbstractBaseModel.
  • Model Registry: Maintains a registry of available models and their corresponding classes.
  • Custom Model Registration: Allows registering custom models by name.

Factory Methods:

Method Description
create(model: Union[str, Model], **kwargs) Creates and returns an instance of a registered model. Accepts either a string representing the model name or a Model enum.
register(model_name: str, model_class: Type[AbstractBaseModel]) Registers a new model to the factory by name. If a model with the same name is already registered, an error is raised.
get_models() Returns a list of all registered model classes.
_register_models() Scans the stylemod.models package and automatically registers all classes inheriting from AbstractBaseModel.

Example Usage:

from stylemod.core.factory import ModelFactory

# Create a model by its enum name (assuming Model.VGG19 is registered)
model = ModelFactory.create("VGG19", content_layer="conv4_2", style_weights={"conv1_1": 1.0})

# Alternatively, create a model by passing a Model enum
from stylemod.models import Model
model = ModelFactory.create(Model.VGG19, content_layer="conv4_2", style_weights={"conv1_1": 1.0})

# Register a custom model
class MyCustomModel(BaseModel):
  ...

ModelFactory.register("MY_CUSTOM_MODEL", MyCustomModel)

# Create an instance of the custom model
custom_model = ModelFactory.create("MY_CUSTOM_MODEL", content_layer='conv4_2', style_weights={'conv1_1': 1.0})

License

stylemod is licensed under the MIT License. See the LICENSE file for details.