tshr-d-dragon/NST

1import cv2 as cv2import numpy as np3import torch4from torchvision import transforms5import os6import matplotlib.pyplot as plt7 8 9from models.definitions.vgg_nets import Vgg16, Vgg19, Vgg16Experimental10 11 12IMAGENET_MEAN_255 = [123.675, 116.28, 103.53]13IMAGENET_STD_NEUTRAL = [1, 1, 1]14 15 16#17# Image manipulation util functions18#19 20def load_image(img_path, target_shape=None):21    if not os.path.exists(img_path):22        raise Exception(f'Path does not exist: {img_path}')23    img = cv.imread(img_path)[:, :, ::-1]  # [:, :, ::-1] converts BGR (opencv format...) into RGB24 25    if target_shape is not None:  # resize section26        if isinstance(target_shape, int) and target_shape != -1:  # scalar -> implicitly setting the height27            current_height, current_width = img.shape[:2]28            new_height = target_shape29            new_width = int(current_width * (new_height / current_height))30            img = cv.resize(img, (new_width, new_height), interpolation=cv.INTER_CUBIC)31        else:  # set both dimensions to target shape32            img = cv.resize(img, (target_shape[1], target_shape[0]), interpolation=cv.INTER_CUBIC)33 34    # this need to go after resizing - otherwise cv.resize will push values outside of [0,1] range35    img = img.astype(np.float32)  # convert from uint8 to float3236    img /= 255.0  # get to [0, 1] range37    return img38 39 40def prepare_img(img_path, target_shape, device):41    img = load_image(img_path, target_shape=target_shape)42 43    # normalize using ImageNet's mean44    # [0, 255] range worked much better for me than [0, 1] range (even though PyTorch models were trained on latter)45    transform = transforms.Compose([46        transforms.ToTensor(),47        transforms.Lambda(lambda x: x.mul(255)),48        transforms.Normalize(mean=IMAGENET_MEAN_255, std=IMAGENET_STD_NEUTRAL)49    ])50 51    img = transform(img).to(device).unsqueeze(0)52 53    return img54 55 56def save_image(img, img_path):57    if len(img.shape) == 2:58        img = np.stack((img,) * 3, axis=-1)59    cv.imwrite(img_path, img[:, :, ::-1])  # [:, :, ::-1] converts rgb into bgr (opencv contraint...)60 61 62def generate_out_img_name(config):63    prefix = os.path.basename(config['content_img_name']).split('.')[0] + '_' + os.path.basename(config['style_img_name']).split('.')[0]64    # called from the reconstruction script65    if 'reconstruct_script' in config:66        suffix = f'_o_{config["optimizer"]}_h_{str(config["height"])}_m_{config["model"]}{config["img_format"][1]}'67    else:68        suffix = f'_o_{config["optimizer"]}_i_{config["init_method"]}_h_{str(config["height"])}_m_{config["model"]}_cw_{config["content_weight"]}_sw_{config["style_weight"]}_tv_{config["tv_weight"]}{config["img_format"][1]}'69    return prefix + suffix70 71 72def save_and_maybe_display(optimizing_img, dump_path, config, img_id, num_of_iterations, should_display=False):73    saving_freq = config['saving_freq']74    out_img = optimizing_img.squeeze(axis=0).to('cpu').detach().numpy()75    out_img = np.moveaxis(out_img, 0, 2)  # swap channel from 1st to 3rd position: ch, _, _ -> _, _, chr76 77    # for saving_freq == -1 save only the final result (otherwise save with frequency saving_freq and save the last pic)78    if img_id == num_of_iterations-1 or (saving_freq > 0 and img_id % saving_freq == 0):79        img_format = config['img_format']80        out_img_name = str(img_id).zfill(img_format[0]) + img_format[1] if saving_freq != -1 else generate_out_img_name(config)81        dump_img = np.copy(out_img)82        dump_img += np.array(IMAGENET_MEAN_255).reshape((1, 1, 3))83        dump_img = np.clip(dump_img, 0, 255).astype('uint8')84        cv.imwrite(os.path.join(dump_path, out_img_name), dump_img[:, :, ::-1])85 86    if should_display:87        plt.imshow(np.uint8(get_uint8_range(out_img)))88        plt.show()89 90 91def get_uint8_range(x):92    if isinstance(x, np.ndarray):93        x -= np.min(x)94        x /= np.max(x)95        x *= 25596        return x97    else:98        raise ValueError(f'Expected numpy array got {type(x)}')99 100 101#102# End of image manipulation util functions103#104 105 106# initially it takes some time for PyTorch to download the models into local cache107def prepare_model(model, device):108    # we are not tuning model weights -> we are only tuning optimizing_img's pixels! (that's why requires_grad=False)109    experimental = False110    if model == 'vgg16':111        if experimental:112            # much more flexible for experimenting with different style representations113            model = Vgg16Experimental(requires_grad=False, show_progress=True)114        else:115            model = Vgg16(requires_grad=False, show_progress=True)116    elif model == 'vgg19':117        model = Vgg19(requires_grad=False, show_progress=True)118    else:119        raise ValueError(f'{model} not supported.')120 121    content_feature_maps_index = model.content_feature_maps_index122    style_feature_maps_indices = model.style_feature_maps_indices123    layer_names = model.layer_names124 125    content_fms_index_name = (content_feature_maps_index, layer_names[content_feature_maps_index])126    style_fms_indices_names = (style_feature_maps_indices, layer_names)127    return model.to(device).eval(), content_fms_index_name, style_fms_indices_names128 129 130def gram_matrix(x, should_normalize=True):131    (b, ch, h, w) = x.size()132    features = x.view(b, ch, w * h)133    features_t = features.transpose(1, 2)134    gram = features.bmm(features_t)135    if should_normalize:136        gram /= ch * h * w137    return gram138 139 140def total_variation(y):141    return torch.sum(torch.abs(y[:, :, :, :-1] - y[:, :, :, 1:])) + \142           torch.sum(torch.abs(y[:, :, :-1, :] - y[:, :, 1:, :]))143