Upload 5 files

app.py

@@ -0,0 +1,80 @@
+import streamlit as st
+import os
+from PIL import Image
+from torchvision.models import vgg19
+from model import StyleTransferModel
+from trainer import trainer_fn
+from utils import process_image, tensor_to_image
+
+
+base_model = vgg19(pretrained=True).features
+final_model = StyleTransferModel(base_model)
+
+# define the title of the app
+st.title('Style Transfer App')
+
+# define the description of the app
+st.write('This app applies the style of one image to another image. This can be used to create artistic images.')
+
+# get all image files in the 'styles' folder
+image_files = [f for f in os.listdir('styles') if f.lower().endswith(('png', 'jpg', 'jpeg', 'gif', 'bmp'))]
+
+# display the images
+st.write('Select style art to apply into your image:')
+
+# Check how many images are available and set columns accordingly
+num_images = len(image_files)
+cols = st.columns(num_images)
+
+# Define the size to which the images will be resized (width, height)
+resize_width = 300
+resize_height = 300
+
+# show each image in a corresponding column
+for idx, img_file in enumerate(image_files):
+    with cols[idx]:
+        st.write(f"Style {idx + 1}")
+        img_path = f'styles/{img_file}'
+        img = Image.open(img_path)
+        
+        # Redimensionar la imagen
+        img_resized = img.resize((resize_width, resize_height))
+        
+        st.image(img_resized, use_container_width=True)
+
+# create a file uploader for the content image
+st.write('Upload the content image:')
+content_image = st.file_uploader('Content Image', type=['png', 'jpg', 'jpeg'])
+
+# create the botton to select the style image between 1, 2, and 3
+choice = st.selectbox('Select the style art:', [f'Style {i + 1}' for i in range(num_images)])
+
+# create a button to run the model
+if st.button('Apply Style Transfer'):
+    if content_image is not None:
+        # get the content image
+        content_img = Image.open(content_image)
+    
+        # get the style image
+        style_choice = choice.split()[-1]  # Extract style number from "Style 1", "Style 2", etc.
+        style_img = Image.open(os.path.join('styles', image_files[int(style_choice) - 1]))  # Get full path
+
+        # preprocess the images
+        content_img = process_image(content_img)
+        style_img = process_image(style_img)
+
+        # run the model
+        st.write('Applying Style Transfer...')
+        target_image = trainer_fn(
+            content_img, style_img, content_img.clone().requires_grad_(True), final_model
+        )
+
+        # convert the tensor to image
+        target_image = tensor_to_image(target_image.squeeze(0))
+
+        # display the result
+        st.write('Result:')
+        st.image(target_image, use_container_width=True)
+    else:
+        st.write('Please upload a content image')
+

loss.py

@@ -0,0 +1,56 @@
+import torch
+import torch.nn as nn
+
+class StyleTransferLoss(nn.Module):
+    def __init__(self, model, content_img, style_img, device="cuda"):
+        super(StyleTransferLoss, self).__init__()
+        self.device = device
+        self.content_img = content_img.to(device)
+        self.style_img = style_img.to(device)
+        self.model = model.to(device)
+
+    def gram_matrix(self, feature_maps):
+        """
+        Calculate Gram Matrix for style features
+        """
+        B, C, H, W = feature_maps.size()
+        features = feature_maps.view(B * C, H * W)
+        G = torch.mm(features, features.t())
+        # Normalize by total elements
+        return G.div(B * C * H * W)
+
+    def get_features(self, image):
+        """
+        Get content and style features from the image
+        """
+        return self.model(image)
+
+    def content_loss(self, target_features, content_features):
+        """
+        Calculate content loss between target and content features
+        """
+        return torch.mean((target_features - content_features) ** 2)
+
+    def style_loss(self, target_features, style_features):
+        """
+        Calculate style loss between target and style features
+        """
+        loss = 0.0
+        for key in self.model.style_layers:
+            target_gram = self.gram_matrix(target_features[key])
+            style_gram = self.gram_matrix(style_features[key])
+            loss += torch.mean((target_gram - style_gram) ** 2)
+        return loss
+
+    def total_loss(
+        self, target_features, content_features, style_features, alpha=1, beta=1e8
+    ):
+        """
+        Calculate total loss (weighted sum of content and style losses)
+        """
+        content = self.content_loss(
+            target_features["block4"], content_features["block4"]
+        )
+        style = self.style_loss(target_features, style_features)
+
+        return alpha * content + beta * style

model.py

@@ -0,0 +1,60 @@
+import torch.nn as nn
+
+class StyleTransferModel(nn.Module):
+    def __init__(self, base_model):
+        super(StyleTransferModel, self).__init__()
+        vgg19 = base_model
+        # Freeze the parameters
+        for param in vgg19.parameters():
+            param.requires_grad = False
+
+        # Split VGG19 into blocks for feature extraction
+        self.block1 = vgg19[:4]  # conv1_1, relu, conv1_2, relu
+        self.pool1 = vgg19[4]  # maxpool
+        self.block2 = vgg19[5:9]  # conv2_1, relu, conv2_2, relu
+        self.pool2 = vgg19[9]  # maxpool
+        self.block3 = vgg19[10:18]  # conv3_1 to relu3_4
+        self.pool3 = vgg19[18]  # maxpool
+        self.block4 = vgg19[19:27]  # conv4_1 to relu4_4
+        self.pool4 = vgg19[27]  # maxpool
+        self.block5 = vgg19[28:36]  # conv5_1 to relu5_4
+
+        # Define content and style layers
+        self.content_layers = ["block4"]  # We'll use output of block4 for content
+        self.style_layers = [
+            "block1",
+            "block2",
+            "block3",
+            "block4",
+            "block5",
+        ]  # All blocks for style
+
+    def forward(self, x):
+        # create a dict to save the results
+        features = {}
+
+        # Block 1
+        x = self.block1(x)
+        features["block1"] = x
+        x = self.pool1(x)
+
+        # Block 2
+        x = self.block2(x)
+        features["block2"] = x
+        x = self.pool2(x)
+
+        # Block 3
+        x = self.block3(x)
+        features["block3"] = x
+        x = self.pool3(x)
+
+        # Block 4
+        x = self.block4(x)
+        features["block4"] = x
+        x = self.pool4(x)
+
+        # Block 5
+        x = self.block5(x)
+        features["block5"] = x
+
+        return features

trainer.py

@@ -0,0 +1,32 @@
+from torch.optim import Adam
+from loss import StyleTransferLoss
+import torch
+
+def trainer_fn(content, style, target_image, model):
+    optimizer = Adam([target_image], lr=0.1)
+    loss_fn = StyleTransferLoss(
+        model=model, content_img=content, style_img=style, device="cpu"
+    )
+
+    with torch.no_grad():
+        content_features = loss_fn.get_features(content.to("cpu"))
+        style_features = loss_fn.get_features(style.to("cpu"))
+    
+    EPOCHS = 100
+    for epoch in range(EPOCHS):
+        # set the gradients to zero
+        optimizer.zero_grad()
+
+        # get the features of the target image
+        target_features = loss_fn.get_features(target_image)
+
+        # calculate the total loss
+        loss = loss_fn.total_loss(target_features, content_features, style_features)
+
+        # backpropagate
+        loss.backward()
+
+        # update the weights
+        optimizer.step()
+
+    return target_image

utils.py

@@ -0,0 +1,32 @@
+from torchvision.transforms import transforms
+import torch
+
+
+def process_image(image, shape=(500, 500)):
+    """
+    This function takes an image and transforms it into a tensor
+
+    """
+    transform = transforms.Compose(
+                [
+                    transforms.Resize(shape),
+                    transforms.ToTensor(),
+                    transforms.Normalize(
+                        mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
+                    ),
+                ]
+            )
+    image = transform(image).unsqueeze(0)
+    return image
+
+def tensor_to_image(tensor):
+    """
+    This function takes a tensor and transforms it into an image
+    """
+    inverse_normalize = transforms.Normalize(
+            mean=[-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225],
+            std=[1 / 0.229, 1 / 0.224, 1 / 0.225],
+        )
+    tensor = inverse_normalize(tensor)
+    tensor = torch.clamp(tensor, 0, 1)
+    return transforms.ToPILImage()(tensor)