evaluate_style_error.py

import argparse
import os
import sys

from skimage import io, transform
from torch.utils.data import DataLoader
import torch
torch.cuda.empty_cache()

import gc
gc.collect()

import re
import numpy as np
from scipy import misc
from PIL import Image
from torch.autograd import Variable
import glob
from torchvision import transforms
import matplotlib.pyplot as plt
import cv2
from skimage.util import random_noise

import lpips

from vgg import Vgg16
import utils
from skimage.transform import swirl, warp

# https://github.com/safwankdb/ReCoNet-PyTorch/blob/master/testwarp.py

def get_subdirectories(directory_path):
    subdirectories = []
    for entry in os.scandir(directory_path):
        if entry.is_dir():
            subdirectories.append(entry.name)
    return subdirectories

def visualize_image(tensor_image):
    # Convert tensor to PIL Image for visualization
    image = transforms.ToPILImage()(tensor_image)

    # Display the image
    plt.imshow(image)
    plt.show()

def main():
    
    parser = argparse.ArgumentParser(description='parser for evaluating a model')
    parser.add_argument("--stylized", type=str, required=True,
                        help="folder that contains the images")
    parser.add_argument("--cuda", type=int, default=1, required=False,
                                  help="use cuda")
    parser.add_argument("--image-size", type=int, default=512,
                                  help="the image size")
    parser.add_argument("--style-image", type=str, required=False,
                                  help="the style image")
    parser.add_argument("--perturbation", type=float, default=0.1)
    parser.add_argument("--noise", type=int, default=0, help="add noise/perturbation to the image")

    args = parser.parse_args()

    if args.cuda and not torch.cuda.is_available():
        print("ERROR: cuda is not available, try running on CPU")
        sys.exit(1)

     # set up midas
    device = torch.device("cuda" if args.cuda else "cpu")
    print("Device: ", torch.cuda.get_device_name(0))
    print("Running on ", device)

    mse_loss = torch.nn.MSELoss()
    lpips_sim = lpips.LPIPS(net='squeeze').to(device)
    vgg = Vgg16(requires_grad=False).to(device)

    style_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Lambda(lambda x: x.mul(255))
    ])
    # style = utils.load_image(args.style_image)
    # style = style_transform(style).to(device)
    # style = style.repeat(1, 1, 1, 1).to(device)
    
    # features_style = vgg(utils.normalize_batch(style))
    # gram_style = [utils.gram_matrix(y) for y in features_style]

    
    sum_content = 0.
    sum_style = 0.
    #############################################################################

    # benchmark_dir = 'benchmark_unstylized'
    benchmark_dir = 'testing_images/fid_all'
    # retrieve the original frames
    original_frames_path = benchmark_dir + '/*.jpg'
    original_frames = []
    for img in sorted(glob.glob(original_frames_path)):
        # image = Image.open(img).convert('RGB')            
        # image = transforms.Resize((args.image_size,args.image_size))(image)
        # image = transforms.ToTensor()(image)
        # image = image.unsqueeze(0).to(device)
        original_frames.append(img)
    print('original frames: ', len(original_frames))

    # styles_path = benchmark_dir = '../../Unity/arbitrary-nst-test/style/' # 'testing_images/tested_image' #  'styles'
    styles_path = 'testing_images/fid_all'
    # retrieve the original frames
    styles_path = styles_path + '/*.*'
    all_styles = []
    for img in sorted(glob.glob(styles_path)):
        # print(img)
        # image = Image.open(img).convert('RGB')            
        # image = transforms.Resize((args.image_size,args.image_size))(image)
        # image = transforms.ToTensor()(image)
        # image = image.unsqueeze(0).to(device)
        all_styles.append(img)
    print('all styles: ', len(all_styles))

    content_error = 0.
    style_error = 0.
    count = 0
    test_frames_path = args.stylized
    print(test_frames_path)
    stylized_images = []
    for stylized_img in sorted(glob.glob(test_frames_path + '/*.*')): 
        stylized_image = Image.open(stylized_img).convert('RGB')        
        stylized_image = transforms.Resize((args.image_size,args.image_size))(stylized_image)
        stylized_image = transforms.ToTensor()(stylized_image)

        if (args.noise):
            perturbation_strength = args.perturbation
            perturbation = torch.randn(stylized_image.size()) * perturbation_strength
            perturbed_image = stylized_image + perturbation
            perturbed_image = torch.clamp(perturbed_image, 0, 1)
            stylized_image = perturbed_image

            # gaussian blur
            # kernel_sizes = {0.1:5, 0.25:11, 0.5: 17, 1.0:23, 2.0: 29, 3.0:35}
            # gaussian_blur = transforms.GaussianBlur(kernel_size=(kernel_sizes[args.perturbation],kernel_sizes[args.perturbation]), sigma=args.perturbation*3)
            # stylized_image = gaussian_blur(stylized_image)

            # salt and pepper
            # stylized_image = torch.tensor(random_noise(stylized_image, mode='s&p', amount=0.1*args.perturbation, clip=True))

            # gaussian noise
            # stylized_image = torch.tensor(random_noise(stylized_image, mode='gaussian', mean=0, var=0.1*args.perturbation, clip=True))
            # stylized_image = stylized_image.float()

            # implanted black rectangles
            # stylized_image = utils.add_black_rectangles_tensor(stylized_image.clone(), 10+int(10*args.perturbation), (10+int(30*args.perturbation), 10+int(30*args.perturbation)))

            # swirl
            # image_np = stylized_image.permute(1, 2, 0).numpy()
            # swirled = swirl(image_np, radius=100 + (140 * args.perturbation), rotation=0, strength=1+int(args.perturbation), clip=True, center=(image_np.shape[1] / 2, image_np.shape[0] / 2), order=1)
            # stylized_image = torch.tensor(swirled).permute(2,0,1)


            # visualize_image(stylized_image)

        stylized_image = stylized_image.unsqueeze(0).to(device)
        stylized_images.append(stylized_image)

        
        # get the original content image
        # print(stylized_img)
        # for content_img in original_frames:
        for style in all_styles:
            base_filename = os.path.basename(style)
            name_without_extension, _ = os.path.splitext(base_filename)
            desired_part = name_without_extension[:-1]
            # print(desired_part)
            if desired_part in stylized_img:
                count += 1
                # print(stylized_img, "   -----    ", desired_part)
                # content_image = Image.open(content_img).convert('RGB')            
                # content_image = transforms.Resize((args.image_size,args.image_size))(content_image)
                # content_image = transforms.ToTensor()(content_image)
                # content_image = content_image.unsqueeze(0).to(device)

                style_image = utils.load_image(style)
                style_image = style_transform(style_image).to(device)
                style_image = style_image.repeat(1, 1, 1, 1).to(device)


                # features_org = vgg(content_image)
                features_stylized = vgg(stylized_image)
                # content_error += mse_loss(features_stylized.relu2_2, features_org.relu2_2)

                features_style = vgg(utils.normalize_batch(style_image))
                gram_style = [utils.gram_matrix(y) for y in features_style]

    
                style_loss = 0.
                for ft_y, gm_s in zip(features_stylized, gram_style):
                    gm_y = utils.gram_matrix(ft_y)
                    style_loss += mse_loss(gm_y, gm_s)
                style_error += style_loss

                print('***************************************************')
                print('Image: ', desired_part, ' Style error: {:.4f}'.format(style_loss * 100))
                break

    print('--------------------------------------------------------------')  
    print('count: ', count, 'Length: ', len(stylized_images))
    # print('Average content error: ', content_error.item()/len(stylized_images))
    print('Average style error: {:.7f}'.format(style_error.item()/len(stylized_images)), '  (Perturbation: {:.2f})'.format(args.perturbation))
    print('--------------------------------------------------------------')
    sum_content += content_error/len(stylized_images)
    sum_style += style_error/len(stylized_images)



    # content_error = 0.
    # style_error = 0.
    # for itr, (org, stylized) in enumerate(zip(original_frames, stylized_images)):
    #     print(stylized.name)
    #     features_org = vgg(org)
    #     features_stylized = vgg(stylized)
    #     content_error += mse_loss(features_stylized.relu2_2, features_org.relu2_2)

    #     style_loss = 0.
    #     for ft_y, gm_s in zip(features_stylized, gram_style):
    #         gm_y = utils.gram_matrix(ft_y)
    #         style_loss += mse_loss(gm_y, gm_s)
    #     style_error += style_loss

    # print('Average content error: ', content_error.item()/len(original_frames))
    # print('Average style error: ', style_error.item()/len(original_frames))
    # sum_content += content_error/len(original_frames)
    # sum_style += style_error/len(original_frames)
        
    
    # print('Average content error over all directories: ', round(sum_content.item()/len(INPUT_DIRS),4))
    # print('Average style error over all directories: ', round(sum_style.item()/len(INPUT_DIRS),8))
        

    
if __name__ == "__main__":
    main()