Spaces:
Runtime error
Runtime error
File size: 8,080 Bytes
9994352 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 |
import os
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import torch
from torchvision import datasets
from torch.utils.data import Dataset, DataLoader, Sampler
import torchvision.transforms as T
# ImageNet Mean and Std values used for Normalization.
Imagenet_Mean = np.array([0.485, 0.456, 0.406])
Imagenet_Std = np.array([0.229, 0.224, 0.225])
def process_img(image_path: os.path, batch_size: int = 1, target_shape = None) -> torch.Tensor:
"""
Function to Pre-process image.
Args:
image_path -> Path of image.
batch_size -> To create batch of images.
target_shape -> Target size of output image.
"""
if not os.path.exists(image_path):
raise Exception(f"Path doesn't exist: {image_path}")
image = Image.open(image_path)
# image = cv2.imread(image_path)[:, :, ::-1]
if target_shape is not None:
if isinstance(target_shape, int) and target_shape != -1:
width, height = image.size
# height, width = image.shape[:2]
new_width = target_shape
new_height = int(height * (new_width / width))
image = image.resize((new_width, new_height))
# image = cv2.resize(image, (new_width, new_height), interpolation = cv2.INTER_CUBIC)
else:
image = image.resize((target_shape[0], target_shape[1]))
# image = cv2.resize(image, (target_shape[0], target_shape[1]), interpolation = cv2.INTER_CUBIC)
image = np.array(image).astype(np.float32)
# image = image.astype(np.float32)
image /= 255.0
transforms = T.Compose([T.ToTensor(),
T.Normalize(mean = Imagenet_Mean, std = Imagenet_Std)
])
image = transforms(image)
image = image.repeat(batch_size, 1, 1, 1)
return image
"""________________________________________________________________________________________________________________________________________________________________"""
"""Class to create a subset of data to be used for training."""
class SequentialSampler(Sampler):
def __init__(self, datasource, subset_size: int = None):
"""
Args:
datasource -> Torch Dataset.
subset_size -> Number of samples to be used.
"""
assert isinstance(datasource, Dataset) or isinstance(datasource, datasets.ImageFolder), "Datasource must be either Torch Dataset or Torchvision ImageFolder."
self.data_source = datasource
if subset_size is None:
subset_size = len(self.data_source)
assert 0 < subset_size <= len(self.data_source), f"Subset size must be between (0, {len(self.data_source)})."
self.subset_size = subset_size
def __iter__(self):
return iter(range(self.subset_size))
def __len__(self):
return self.subset_size
"""________________________________________________________________________________________________________________________________________________________________"""
def Train_DataLoader(training_config: dict) -> DataLoader:
"""
Function to create DataLoader for training.
Args:
training_config -> Dictionary of params config.
"""
transforms = T.Compose([T.Resize(training_config['image_size']),
T.CenterCrop(training_config['image_size']),
T.ToTensor(),
T.Normalize(mean = Imagenet_Mean, std = Imagenet_Std)
])
Train_dataset = datasets.ImageFolder(training_config['dataset_path'], transform = transforms)
sampler = SequentialSampler(Train_dataset, subset_size = training_config['subset_size'])
training_config['subset_size'] = len(sampler)
train_dl = DataLoader(Train_dataset, batch_size = training_config['batch_size'], sampler = sampler, drop_last = True)
return train_dl
"""________________________________________________________________________________________________________________________________________________________________"""
def gram_matrix(feature_map: torch.Tensor, should_mormalize: bool = True):
"""
Function to create Gram matrix.
Gram matrix - It captures the texture information from the feature maps extracted from convolutional
neural network. It is obtained using dot product between the feature maps. In other words,
it is a simple covariance matrix between different feature maps.
Args:
feature_map -> Feature extraced from pre-trained vision models.
should_normalize -> Should normalize the matrix.
"""
(b, ch, h, w) = feature_map.size()
features = feature_map.view(b, ch, h * w)
features_t = features.transpose(1, 2)
gram_mat = features.bmm(features_t)
if should_mormalize:
gram_mat = gram_mat / (ch * h * w)
return gram_mat
"""________________________________________________________________________________________________________________________________________________________________"""
def total_variation_loss(image_batch) -> torch.Tensor:
"""
Function for Total variation loss.
Used for spatial continuity between the pixels of the generated image, thereby denoising it and giving it visual coherence.
Args:
image_batch -> Batch of image tensors.
"""
batch_size = image_batch.shape[0]
tv_height = torch.sum(torch.abs(image_batch[:, :, :-1, :] - image_batch[:, :, 1:, :]))
tv_width = torch.sum(torch.abs(image_batch[:, :, :, :-1]) - image_batch[:, :, :, 1:])
return (tv_height + tv_width) / batch_size
"""________________________________________________________________________________________________________________________________________________________________"""
def post_process_img(image: np.ndarray) -> np.ndarray:
"""
Function to post process the generated image.
Args:
image -> Numpy array of generated image.
"""
assert isinstance(image, np.ndarray), f"Expected Numpy Array, but got {type(image)}"
mean = Imagenet_Mean.reshape(-1, 1, 1)
std = Imagenet_Std.reshape(-1, 1, 1)
image = (image * std) + mean
image = (np.clip(image, 0., 1.) * 255).astype(np.uint8)
image = np.moveaxis(image, 0, 2)
return image
"""________________________________________________________________________________________________________________________________________________________________"""
def save_and_display(config: dict, image: np.ndarray, should_display: bool = False, index: int = None):
"""
Function to save and display images.
Args:
config -> Dictionary of config.
image -> Numpy array of image.
should_display -> To display image.
index -> Index of generated image during training.
"""
assert isinstance(image, np.ndarray), f"Expected Numpy Array, but got {type(image)}"
image = post_process_img(image)
if index is not None:
fake_fname = 'generated-images-{0:0=4d}.jpg'.format(index)
else:
fake_fname = f'Stylized-image-{config["content_img_name"].split(".")[0]}.jpg'
os.makedirs(config['save_folder'], exist_ok = True)
plt.imsave(os.path.join(config['save_folder'], fake_fname), image)
# cv2.imwrite(os.path.join(config['save_folder'], fake_fname), image[:, :, ::-1])
if should_display:
plt.imshow(image)
plt.show()
"""________________________________________________________________________________________________________________________________________________________________"""
def print_parameters(model: torch.nn.Module) -> None:
"""
Function to print number of parameters.
Args:
model -> Torch model
"""
Num_of_parameters = sum(p.numel() for p in model.parameters())
print("Model Parameters : {:.3f} M".format(Num_of_parameters / 1e6)) |