Image doesn't save on using .show() (PIL, Style Transfer)

28 Views Asked by At

I'm sending a POST request to my flask API endpoint, which has the ML model for image style transfer. On sending the request, I do get an array representation of the image, but the image doesn't actually get saved on my Mac.

The POST request using curl:


    curl -X POST -F "[email protected]" -F "[email protected]" http://127.0.0.1:5000/style_transfer

POST request using client test.py:

    import requests
    from PIL import Image
    import numpy as np

    # Replace with the actual file paths of your content and style images
    content_image_path = 'octopus.jpg'
    style_image_path = 'hockney.jpg'

    # Upload the images to the API
    files = {'content_image': ('content_image.jpg', open(content_image_path, 'rb')),
         'style_image': ('style_image.jpg', open(style_image_path, 'rb'))}

    response = requests.post('http://127.0.0.1:5000/style_transfer', files=files)

    # Check if the request was successful
    if response.status_code == 200:
        result = response.json()
        resultant_image_data = result['resultant_image']

        # Convert the resultant image data to a NumPy array
        resultant_image_np = np.array(resultant_image_data)*255.0
        resultant_image_np = resultant_image_data.astype(np.uint8)
    
        # Create a PIL Image from the NumPy array
        resultant_image = Image.fromarray(resultant_image_np, "RGB")
    
        # Save the image to a file
        resultant_image.show()
    
        resultant_image.save('output.jpg') 
    else:
        print('Error:', response.status_code, response.text)

My API (app.py) code:

    from flask import Flask, request, jsonify
    import numpy as np
    mport torch
    from torchvision import transforms, models
    from PIL import Image

    app = Flask(__name__)

    # Load the pre-trained VGG model as before
    vgg = models.vgg19(pretrained=True).features

    # Freeze all VGG parameters since we're only optimizing the target image
    for param in vgg.parameters():
        param.requires_grad_(False)
    
    # Move the model to GPU, if available
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    vgg.to(device)

    # Helper function to preprocess an image
    def transform_image(image):
        # Resize and normalize the image
        in_transform = transforms.Compose([
            transforms.Resize(400),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
        ])
        image = in_transform(image).unsqueeze(0).to(device)
        return image

    # Helper function to convert a tensor to a NumPy image
    def im_convert(tensor):
        image = tensor.to("cpu").clone().detach()
        image = image.numpy().squeeze()
        image = image.transpose(1, 2, 0)
        image = image * np.array((0.229, 0.224, 0.225)) + np.array((0.485, 0.456, 0.406))
        image = image.clip(0, 1)
        return image

    # Perform style transfer
    def perform_style_transfer(content_image, style_image, steps=1):
        # Load and preprocess the content and style images
        content_image = transform_image(content_image)
        style_image = transform_image(style_image)

        # Define the target image as a copy of the content image
        # target = content_image.clone().requires_grad_(True)

        def get_features(image, model, layers=None):
            """ Run an image forward through a model and get the features for 
                a set of layers. Default layers are for VGGNet matching Gatys et al (2016)
            """
        
            ## TODO: Complete mapping layer names of PyTorch's VGGNet to names from the paper
            ## Need the layers for the content and style representations of an image
            if layers is None:
                layers = {'0': 'conv1_1',
                        '5': 'conv2_1', 
                        '10': 'conv3_1', 
                        '19': 'conv4_1',
                        '21': 'conv4_2',  ## content representation
                        '28': 'conv5_1'}
            
            features = {}
            x = image
            # model._modules is a dictionary holding each module in the model
            for name, layer in model._modules.items():
                x = layer(x)
                if name in layers:
                    features[layers[name]] = x
                
            return features

        def gram_matrix(tensor):
            """ Calculate the Gram Matrix of a given tensor 
                Gram Matrix: https://en.wikipedia.org/wiki/Gramian_matrix
            """
    
            # get the batch_size, depth, height, and width of the Tensor
            b, d, h, w = tensor.size()
        
            # reshape so we're multiplying the features for each channel
            tensor = tensor.view(b * d, h * w)
                
            # calculate the gram matrix
            gram = torch.mm(tensor, tensor.t())
        
            return gram 

        # get content and style features only once before training
        content_features = get_features(content_image, vgg)
        style_features = get_features(style_image, vgg)

        # calculate the gram matrices for each layer of our style representation
        style_grams = {layer: gram_matrix(style_features[layer]) for layer in style_features}

        # create a third "target" image and prep it for change
        # it is a good idea to start off with the target as a copy of our *content* image
        # then iteratively change its style
        target = content_image.clone().requires_grad_(True).to(device)

        # Define style weights and other parameters (similar to the original code)
        style_weights = {'conv1_1': 1.0, 'conv2_1': 0.75, 'conv3_1': 0.2, 'conv4_1': 0.2, 'conv5_1': 0.2}
        content_weight = 1  # alpha
        style_weight = 1e6  # beta

        optimizer = torch.optim.Adam([target], lr=0.003)

        for ii in range(1, steps+1):
        
            # get the features from your target image
            target_features = get_features(target, vgg)
        
            # the content loss
            content_loss = torch.mean((target_features['conv4_2'] - content_features['conv4_2'])**2)
        
            # the style loss
            # initialize the style loss to 0
            style_loss = 0
            # then add to it for each layer's gram matrix loss
            for layer in style_weights:
                # get the "target" style representation for the layer
                target_feature = target_features[layer]
                target_gram = gram_matrix(target_feature)
                _, d, h, w = target_feature.shape
                # get the "style" style representation
                style_gram = style_grams[layer]
                # the style loss for one layer, weighted appropriately
                layer_style_loss = style_weights[layer] * torch.mean((target_gram - style_gram)**2)
                # add to the style loss
                style_loss += layer_style_loss / (d * h * w)
            
            # calculate the *total* loss
            total_loss = content_weight * content_loss + style_weight * style_loss
        
            # update your target image
            optimizer.zero_grad()
            total_loss.backward()
            optimizer.step()
        # Get the resultant image as a NumPy array
        resultant_image = im_convert(target)

        return resultant_image
    
    # Define the route to accept POST requests with two images
    @app.route('/style_transfer', methods=['POST'])
    def style_transfer():
        try:
            # Get the uploaded images from the request
            content_image = request.files['content_image']
            style_image = request.files['style_image']

            # Perform style transfer
            resultant_image = perform_style_transfer(
                Image.open(content_image),  # Convert FileStorage to Image
                Image.open(style_image)     # Convert FileStorage to Image
            )

            # Create a response dictionary with the resultant image
            response_data = {
                'resultant_image': resultant_image.tolist()  # Convert to a list for JSON serialization
            }

            return jsonify(response_data)
        except Exception as e:
            return jsonify({'error': str(e)})

    if __name__ == '__main__':
        app.run(debug=True)

All the images are in the same directory, so I've specified only the name as the path. I tried specifying all types of paths, but nothing changed. I've kept the value of 'step' variable 1 so that it doesn't take too long to complete the POST request. I tried POST request using curl and a client. Neither worked.

The output I get on making POST request:

    {
      "resultant_image": [  
        [
         [
            0.1999999632835388,
            0.18039219665527345,
            0.10196077203750614
          ],
         [
           ... //many such numbers
          ]
        ]
      ]
    }
0

There are 0 best solutions below