first commit

.gitignore

@@ -0,0 +1,133 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+.github
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Lightning /research  
+test_tube_exp/
+tests/tests_tt_dir/
+tests/save_dir
+default/   
+data/
+test_tube_logs/
+test_tube_data/
+datasets/
+model_weights/
+tests/save_dir  
+tests/tests_tt_dir/
+processed/
+raw/
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+
+# IDEs
+.idea
+.vscode
+
+# seed project
+lightning_logs/
+MNIST
+.DS_Store
+
+# ignore CI testing for now
+.github/ 
+tests/

gradio_cached_examples/14/log.csv

@@ -0,0 +1,2 @@
+output,flag,username,timestamp
+/Users/abhiroop/Developer/neural-style-transfer/gradio_cached_examples/14/output/b3bac432e10f60f0fc74d6dd737f3e5b190aff2f/tmpftrz8no9.png,,,2023-09-07 19:03:10.521984

nst/losses.py

@@ -0,0 +1,42 @@
+import torch 
+from torch import nn 
+import torch.nn.functional as F 
+
+class ContentLoss(nn.Module):
+    """
+    Content Loss for the neural style transfer algorithm.
+    """
+    def __init__(self, target: torch.Tensor, device: torch.device) -> None:
+        super(ContentLoss, self).__init__()
+        batch_size, channels, height, width = target.size()
+        target = target.view(batch_size * channels, height * width)
+        self.target = target.detach().to(device)
+
+    def __str__(self) -> str:
+        return "Content loss"
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        batch_size, channels, height, width = input.size()
+        input = input.view(batch_size * channels, height * width)
+        return F.mse_loss(input, self.target)
+
+
+class StyleLoss(nn.Module):
+    """
+    Style loss for the neural style transfer algorithm.
+    """
+    def __init__(self, target: torch.Tensor, device: torch.device) -> None:
+        super(StyleLoss, self).__init__()
+        self.target = self.compute_gram_matrix(target).detach().to(device)
+
+    def __str__(self) -> str:
+        return "Style loss"
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:
+        input = self.compute_gram_matrix(input)
+        return F.mse_loss(input, self.target)
+
+    def compute_gram_matrix(self, input: torch.Tensor) -> torch.Tensor:
+        batch_size, channels, height, width = input.size()
+        input = input.view(batch_size * channels, height * width)
+        return torch.matmul(input, input.T).div(batch_size * channels * height * width)

nst/models/vgg19.py

@@ -0,0 +1,94 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from typing import Tuple, Dict, Optional
+
+class Normalization(nn.Module):
+    """
+    Normalization module for VGG19.
+    """
+    def __init__(self, mean: torch.Tensor, std: torch.Tensor) -> None:
+        super(Normalization, self).__init__()
+        self.mean = mean.view(-1, 1, 1)
+        self.std = std.view(-1, 1, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return (x - self.mean) / self.std
+
+
+class ConvBlock1(nn.Module):
+    """
+    Convolution block for VGG19 [conv2d, conv2d, maxpool2d].
+    """
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super(ConvBlock1, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, padding=1)
+        self.relu1 = nn.ReLU(inplace=False)
+        self.conv2 = nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, padding=1)
+        self.relu2 = nn.ReLU(inplace=False)
+        self.max_pool2d = nn.MaxPool2d(kernel_size=2)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor]:
+        conv1 = self.relu1(self.conv1(x))
+        conv2 = self.relu2(self.conv2(conv1))
+        out = self.max_pool2d(conv2)
+        return conv1, conv2, out
+
+
+class ConvBlock2(nn.Module):
+    """
+    Convolution block for VGG19 [conv2d, conv2d, conv2d, conv2d, maxpool2d].
+    """
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        super(ConvBlock2, self).__init__()
+        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, padding=1)
+        self.relu1 = nn.ReLU(inplace=False)
+        self.conv2 = nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, padding=1)
+        self.relu2 = nn.ReLU(inplace=False)
+        self.conv3 = nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, padding=1)
+        self.relu3 = nn.ReLU(inplace=False)
+        self.conv4 = nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=3, padding=1)
+        self.relu4 = nn.ReLU(inplace=False)
+        self.max_pool2d = nn.MaxPool2d(kernel_size=2)
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor]:
+        conv1 = self.relu1(self.conv1(x))
+        conv2 = self.relu2(self.conv2(conv1))
+        conv3 = self.relu3(self.conv3(conv2))
+        conv4 = self.relu4(self.conv4(conv3))
+        out = self.max_pool2d(conv2)
+        return conv1, conv2, conv3, conv4, out
+
+
+class VGG19(nn.Module):
+    """
+    VGG19 module with only the feature extractor.
+    """
+    def __init__(self, in_channels: int=3, out_channels: int=64, 
+                mean: Optional[torch.Tensor]=None, std: Optional[torch.Tensor]=None) -> None:
+        super(VGG19, self).__init__()
+        self.norm = Normalization(mean=mean, std=std)
+        self.conv1 = ConvBlock1(in_channels=in_channels, out_channels=out_channels)
+        self.conv2 = ConvBlock1(in_channels=out_channels, out_channels=out_channels * 2)
+        self.conv3 = ConvBlock2(in_channels=out_channels * 2, out_channels=out_channels * 4)
+        self.conv4 = ConvBlock2(in_channels=out_channels * 4, out_channels=out_channels * 8)
+        self.conv5 = ConvBlock2(in_channels=out_channels * 8, out_channels = out_channels * 8)
+            
+    def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]: 
+        x = self.norm(x)
+        conv1_1, conv1_2, out = self.conv1(x)
+        conv2_1, conv2_2, out = self.conv2(out)
+        conv3_1, conv3_2, conv3_3, conv3_4, out = self.conv3(out)
+        conv4_1, conv4_2, conv4_3, conv4_4, out = self.conv4(out)
+        conv5_1, conv5_2, conv5_3, conv5_4, out = self.conv5(out)
+        outputs = {
+            "conv1": [conv1_1, conv1_2],
+            "conv2": [conv2_1, conv2_2],
+            "conv3": [conv3_1, conv3_2, conv3_3, conv3_4],
+            "conv4": [conv4_1, conv4_2, conv4_3, conv4_4],
+            "conv5": [conv5_1, conv5_2, conv5_3, conv5_4],
+            "out" : out
+        }
+
+        return outputs 

nst/train.py

@@ -0,0 +1,229 @@
+import os 
+import sys
+sys.path.append(os.path.abspath(os.path.pardir))
+from argparse import ArgumentParser
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+import torch.optim as optim
+
+from torchvision import transforms
+from torchvision.models import vgg19
+
+from PIL import Image
+import matplotlib.pyplot as plt
+
+from nst.models.vgg19 import VGG19
+from nst.losses import ContentLoss, StyleLoss
+
+from tqdm import tqdm
+from typing import List, Union
+
+def main() -> None:
+    # command line args 
+    parser = ArgumentParser()
+    parser.add_argument("--use_gpu", default=True, type=bool)
+    parser.add_argument("--content_dir", default="../images/content/dancing.jpg", type=str)
+    parser.add_argument("--style_dir", default="../images/style/picasso.jpg", type=str)
+    parser.add_argument("--input_image", default="content", type=str)
+    parser.add_argument("--output_dir", default="../result/result.jpg", type=str)
+    parser.add_argument("--iterations", default=100, type=int)
+    parser.add_argument("--alpha", default=1, type=int)
+    parser.add_argument("--beta", default=1000000, type=int)
+    parser.add_argument("--style_layer_weight", default=1.0, type=float)
+    args = parser.parse_args()
+
+    device = torch.device("cuda") if (torch.cuda.is_available() and args.use_gpu) else torch.device("cpu")
+    print(f"training on device {device}")
+
+    # content and style images
+    content = image_loader(args.content_dir, device)
+    style = image_loader(args.style_dir, device)
+
+    # input image
+    if args.input_image == "content":
+        x = content.clone()
+    elif args.input_image == "style":
+        x = style.clone()
+    else:
+        x = torch.randn(content.data.size(), device=device)
+
+    # mean and std for vgg19
+    mean = torch.tensor([0.485, 0.456, 0.406]).to(device)
+    std = torch.tensor([0.229, 0.224, 0.225]).to(device)
+
+    # vgg19 model
+    model = VGG19(mean=mean, std=std).to(device=device)
+    model = load_vgg19_weights(model, device)
+    # LBFGS optimizer like in paper
+    optimizer = optim.LBFGS([x.requires_grad_()])
+
+    # computing content and style representations
+    content_outputs = model(content)
+    style_outputs = model(style)
+
+    # defining content and style losses
+    content_loss = ContentLoss(content_outputs["conv4"][1], device)
+    style_losses = []
+    for i in range(1, 6):
+        style_losses.append(StyleLoss(style_outputs[f"conv{i}"][0], device))
+
+    # run style transfer
+    output = train(model, optimizer, content_loss, style_losses, x,
+                   iterations=args.iterations, alpha=args.alpha, beta=args.beta, 
+                   style_weight=args.style_layer_weight)
+    output = output.detach().to("cpu")
+
+    # save result
+    plt.imsave(args.output_dir, output[0].permute(1, 2, 0).numpy())
+
+def image_loader(path: str, device: torch.device=torch.device("cuda")) -> torch.Tensor:
+    """
+    Loads and resizes the image.
+
+    Args:
+        path (str): Path to the image.
+        device (torch.device): device to load the image in.
+
+    Returns:
+        img (torch.Tensor): Loaded image as torch.Tensor.
+    """
+    transform = transforms.Compose([
+                    transforms.Resize((512, 512)),
+                    transforms.ToTensor(),
+                ])
+    img = Image.open(path)
+    img = transform(img)
+    img = img.unsqueeze(0).to(device=device)
+    return img
+
+def load_vgg19_weights(model: nn.Module, device: torch.device) -> nn.Module:
+    """
+    Loads VGG19 pretrained weights from ImageNet for style transfer.
+    
+    Args:
+        model (nn.Module): VGG19 feature module with randomized weights.
+        device (torch.device): The device to load the model in. 
+
+    Returns:
+        model (nn.Module): VGG19 module with pretrained ImageNet weights loaded.
+    """
+    pretrained_model = vgg19(pretrained=True).features.to(device).eval()
+
+    matching_keys = {
+        "conv1.conv1.weight": "0.weight",
+        "conv1.conv1.bias": "0.bias",
+        "conv1.conv2.weight": "2.weight",
+        "conv1.conv2.bias": "2.bias",
+
+        "conv2.conv1.weight": "5.weight",
+        "conv2.conv1.bias": "5.bias",
+        "conv2.conv2.weight": "7.weight",
+        "conv2.conv2.bias": "7.bias",
+
+        "conv3.conv1.weight": "10.weight",
+        "conv3.conv1.bias": "10.bias",
+        "conv3.conv2.weight": "12.weight",
+        "conv3.conv2.bias": "12.bias",
+        "conv3.conv3.weight": "14.weight",
+        "conv3.conv3.bias": "14.bias",
+        "conv3.conv4.weight": "16.weight",
+        "conv3.conv4.bias": "16.bias",
+
+        "conv4.conv1.weight": "19.weight",
+        "conv4.conv1.bias": "19.bias",
+        "conv4.conv2.weight": "21.weight",
+        "conv4.conv2.bias": "21.bias",
+        "conv4.conv3.weight": "23.weight",
+        "conv4.conv3.bias": "23.bias",
+        "conv4.conv4.weight": "25.weight",
+        "conv4.conv4.bias": "25.bias",
+
+        "conv5.conv1.weight": "28.weight",
+        "conv5.conv1.bias": "28.bias",
+        "conv5.conv2.weight": "30.weight",
+        "conv5.conv2.bias": "30.bias",
+        "conv5.conv3.weight": "32.weight",
+        "conv5.conv3.bias": "32.bias",
+        "conv5.conv4.weight": "34.weight",
+        "conv5.conv4.bias": "34.bias",
+    }
+
+    pretrained_dict = pretrained_model.state_dict()
+    model_dict = model.state_dict()
+
+    for key, value in matching_keys.items():
+        model_dict[key] = pretrained_dict[value]
+    
+    model.load_state_dict(model_dict)
+
+    return model
+
+
+def train(model: nn.Module, optimizer: torch.optim, content_loss: ContentLoss, style_losses: List[StyleLoss], 
+          x: torch.Tensor, iterations: int=100, alpha: int=1, beta: int=1000000, style_weight: Union[int, float]=1.0) -> torch.Tensor:
+    """
+    Train the neural style transfer algorithm.
+
+    Args:
+        model (nn.Module): The VGG19 feature extractor for training the style transfer algorithm.
+        optimizer (torch.optim): The optimization module to use.
+        content_loss (ContentLoss): The content loss to preserve the content representation during style transfer.
+        style_losses (List[StyleLoss]): A list of style loss objects to preserve the style representation across 
+                                        different layers during style transfer.
+        x (torch.Tensor): The input image for style transfer.
+        iterations (int): Number of iterations to run.
+        alpha (int): The weight given to content loss while computing the total loss.
+        beta (int): The weight given to style loss while computing the total loss.
+        style_weight Union[int, float]: The weight given to style loss of each layer while computing total style loss.
+
+    Returns:
+        x (torch.Tensor): The input image with the content and style transfered.
+    """
+
+    with tqdm(range(iterations)) as iterations:
+        for iteration in iterations:
+            iterations.set_description(f"Iteration: {iteration}")
+
+            def closure():
+                optimizer.zero_grad()
+
+                # correcting to 0-1 range
+                x.data.clamp_(0, 1)
+                outputs = model(x)
+
+                # input content and style representations
+                content_feature_maps = outputs["conv4"][1]
+                style_feature_maps = []
+                for i in range(1, 6):
+                    style_feature_maps.append(outputs[f"conv{i}"][0])
+
+                # input content and style losses
+                total_content_loss = content_loss(content_feature_maps)
+
+                total_style_loss = 0
+                for feature_map, style_loss in zip(style_feature_maps, style_losses):
+                    total_style_loss += (style_weight * style_loss(feature_map))
+
+                # total loss
+                loss = (alpha * total_content_loss) + (beta * total_style_loss)
+                loss.backward()
+
+                iterations.set_postfix({
+                    "content loss": total_content_loss.item(),
+                    "style loss": total_style_loss.item(),
+                    "total loss": loss.item()
+                })
+
+                return loss
+
+            optimizer.step(closure)
+
+    # final correction
+    x.data.clamp_(0, 1)
+    return x
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,7 @@
+torch
+torchvision
+numpy
+matplotlib
+tqdm
+pillow
+gradio

setup.cfg

@@ -0,0 +1,43 @@
+[tool:pytest]
+norecursedirs =
+    .git
+    dist
+    build
+addopts =
+    --strict
+    --doctest-modules
+    --durations=0
+
+[coverage:report]
+exclude_lines =
+    pragma: no-cover
+    pass
+
+[flake8]
+max-line-length = 120
+exclude = .tox,*.egg,build,temp
+select = E,W,F
+doctests = True
+verbose = 2
+# https://pep8.readthedocs.io/en/latest/intro.html#error-codes
+format = pylint
+# see: https://www.flake8rules.com/
+ignore =
+    E731  # Do not assign a lambda expression, use a def
+    W504  # Line break occurred after a binary operator
+    F401  # Module imported but unused
+    F841  # Local variable name is assigned to but never used
+    W605  # Invalid escape sequence 'x'
+
+# setup.cfg or tox.ini
+[check-manifest]
+ignore =
+    *.yml
+    .github
+    .github/*
+
+[metadata]
+license_file = LICENSE
+description-file = README.md
+# long_description = file:README.md
+# long_description_content_type = text/markdown

setup.py

@@ -0,0 +1,14 @@
+#!/usr/bin/env python
+
+from setuptools import setup, find_packages
+
+setup(
+    name='nst',
+    version='0.1.0',
+    description='Implementation of Neural Style Transfer using PyTorch',
+    author='Abhiroop Tejomay',
+    author_email='abhirooptejomay@gmail.com',
+    url='https://github.com/visualCalculus/neural-style-transfer',
+    packages=find_packages(include=["nst", "nst.*"]),
+)
+