Upload 8 files

utils/__init__.py

@@ -0,0 +1,21 @@
+from .common import *
+from .image_processing import *
+
+class DefaultArgs:
+    dataset ='Hayao'
+    data_dir ='/content'
+    epochs = 10
+    batch_size = 1
+    checkpoint_dir ='/content/checkpoints'
+    save_image_dir ='/content/images'
+    display_image =True
+    save_interval =2
+    debug_samples =0
+    lr_g = 0.001
+    lr_d = 0.002
+    wadvg = 300.0
+    wadvd = 300.0
+    wcon = 1.5
+    wgra = 3
+    wcol = 10
+    use_sn = False

utils/common.py

@@ -0,0 +1,165 @@
+import torch
+import gc
+import os
+import torch.nn as nn
+import urllib.request
+import cv2
+from tqdm import tqdm
+
+HTTP_PREFIXES = [
+    'http',
+    'data:image/jpeg',
+]
+
+
+RELEASED_WEIGHTS = {
+    "hayao:v2": (
+        # Dataset trained on Google Landmark micro as training real photo
+        "v2",
+        "https://github.com/ptran1203/pytorch-animeGAN/releases/download/v1.1/GeneratorV2_gldv2_Hayao.pt"
+    ),
+    "hayao:v1": (
+        "v1",
+        "https://github.com/ptran1203/pytorch-animeGAN/releases/download/v1.0/generator_hayao.pth"
+    ),
+    "hayao": (
+        "v1",
+        "https://github.com/ptran1203/pytorch-animeGAN/releases/download/v1.0/generator_hayao.pth"
+    ),
+    "shinkai:v1": (
+        "v1",
+        "https://github.com/ptran1203/pytorch-animeGAN/releases/download/v1.0/generator_shinkai.pth"
+    ),
+    "shinkai": (
+        "v1",
+        "https://github.com/ptran1203/pytorch-animeGAN/releases/download/v1.0/generator_shinkai.pth"
+    ),
+}
+
+def is_image_file(path):
+    _, ext = os.path.splitext(path)
+    return ext.lower() in (".png", ".jpg", ".jpeg")
+
+
+def read_image(path):
+    """
+    Read image from given path
+    """
+
+    if any(path.startswith(p) for p in HTTP_PREFIXES):
+        urllib.request.urlretrieve(path, "temp.jpg")
+        path = "temp.jpg"
+
+    return cv2.imread(path)[: ,: ,::-1]
+
+
+def save_checkpoint(model, path, optimizer=None, epoch=None):
+    checkpoint = {
+        'model_state_dict': model.state_dict(),
+        'epoch': epoch,
+    }
+    if optimizer is  not None:
+        checkpoint['optimizer_state_dict'] = optimizer.state_dict()
+
+    torch.save(checkpoint, path)
+
+def maybe_remove_module(state_dict):
+    # Remove added module ins state_dict in ddp training
+    # https://discuss.pytorch.org/t/why-are-state-dict-keys-getting-prepended-with-the-string-module/104627/3
+    new_state_dict = {}
+    module_str = 'module.'
+    for k, v in state_dict.items():
+
+        if k.startswith(module_str):
+            k = k[len(module_str):]
+        new_state_dict[k] = v
+    return new_state_dict
+
+
+def load_checkpoint(model, path, optimizer=None, strip_optimizer=False, map_location=None) -> int:
+    state_dict = load_state_dict(path, map_location)
+    model_state_dict = maybe_remove_module(state_dict['model_state_dict'])
+    model.load_state_dict(
+        model_state_dict,
+        strict=True
+    )
+    if 'optimizer_state_dict' in state_dict:
+        if optimizer is not None:
+            optimizer.load_state_dict(state_dict['optimizer_state_dict'])
+        if strip_optimizer:
+            del state_dict["optimizer_state_dict"]
+            torch.save(state_dict, path)
+            print(f"Optimizer stripped and saved to {path}")
+
+    epoch = state_dict.get('epoch', 0)
+    return epoch
+
+
+def load_state_dict(weight, map_location) -> dict:
+    if weight.lower() in RELEASED_WEIGHTS:
+        weight = _download_weight(weight.lower())
+
+    if map_location is None:
+        # auto select
+        map_location = 'cuda' if torch.cuda.is_available() else 'cpu'
+    state_dict = torch.load(weight, map_location=map_location)
+
+    return state_dict
+
+
+def initialize_weights(net):
+    for m in net.modules():
+        try:
+            if isinstance(m, nn.Conv2d):
+                # m.weight.data.normal_(0, 0.02)
+                torch.nn.init.xavier_uniform_(m.weight)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.ConvTranspose2d):
+                # m.weight.data.normal_(0, 0.02)
+                torch.nn.init.xavier_uniform_(m.weight)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.Linear):
+                # m.weight.data.normal_(0, 0.02)
+                torch.nn.init.xavier_uniform_(m.weight)
+                m.bias.data.zero_()
+            elif isinstance(m, nn.BatchNorm2d):
+                m.weight.data.fill_(1)
+                m.bias.data.zero_()
+        except Exception as e:
+            # print(f'SKip layer {m}, {e}')
+            pass
+
+
+def set_lr(optimizer, lr):
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+
+class DownloadProgressBar(tqdm):
+    '''
+    https://stackoverflow.com/questions/15644964/python-progress-bar-and-downloads
+    '''
+    def update_to(self, b=1, bsize=1, tsize=None):
+        if tsize is not None:
+            self.total = tsize
+        self.update(b * bsize - self.n)
+
+
+def _download_weight(weight):
+    '''
+    Download weight and save to local file
+    '''
+    os.makedirs('.cache', exist_ok=True)
+    url = RELEASED_WEIGHTS[weight][1]
+    filename = os.path.basename(url)
+    save_path = f'.cache/{filename}'
+
+    if os.path.isfile(save_path):
+        return save_path
+
+    desc = f'Downloading {url} to {save_path}'
+    with DownloadProgressBar(unit='B', unit_scale=True, miniters=1, desc=desc) as t:
+        urllib.request.urlretrieve(url, save_path, reporthook=t.update_to)
+
+    return save_path
+

utils/fast_numpyio.py

@@ -0,0 +1,43 @@
+# code from https://github.com/divideconcept/fastnumpyio/blob/main/fastnumpyio.py
+
+import sys
+import numpy as np
+import numpy.lib.format
+import struct
+
+def save(file, array):
+    magic_string=b"\x93NUMPY\x01\x00v\x00"
+    header=bytes(("{'descr': '"+array.dtype.descr[0][1]+"', 'fortran_order': False, 'shape': "+str(array.shape)+", }").ljust(127-len(magic_string))+"\n",'utf-8')
+    if type(file) == str:
+        file=open(file,"wb")
+    file.write(magic_string)
+    file.write(header)
+    file.write(array.data)
+
+def pack(array):
+    size=len(array.shape)
+    return bytes(array.dtype.byteorder.replace('=','<' if sys.byteorder == 'little' else '>')+array.dtype.kind,'utf-8')+array.dtype.itemsize.to_bytes(1,byteorder='little')+struct.pack(f'<B{size}I',size,*array.shape)+array.data
+
+def load(file):
+    if type(file) == str:
+        file=open(file,"rb")
+    header = file.read(128)
+    if not header:
+        return None
+    descr = str(header[19:25], 'utf-8').replace("'","").replace(" ","")
+    shape = tuple(int(num) for num in str(header[60:120], 'utf-8').replace(', }', '').replace('(', '').replace(')', '').split(','))
+    datasize = numpy.lib.format.descr_to_dtype(descr).itemsize
+    for dimension in shape:
+        datasize *= dimension
+    return np.ndarray(shape, dtype=descr, buffer=file.read(datasize))
+
+def unpack(data):
+    dtype = str(data[:2],'utf-8')
+    dtype += str(data[2])
+    size = data[3]
+    shape = struct.unpack_from(f'<{size}I', data, 4)
+    datasize=data[2]
+    for dimension in shape:
+        datasize *= dimension
+    return np.ndarray(shape, dtype=dtype, buffer=data[4+size*4:4+size*4+datasize])
+

utils/image_processing.py

@@ -0,0 +1,135 @@
+import torch
+import cv2
+import os
+import numpy as np
+from tqdm import tqdm
+
+
+def gram(input):
+    """
+    Calculate Gram Matrix
+
+    https://pytorch.org/tutorials/advanced/neural_style_tutorial.html#style-loss
+    """
+    b, c, w, h = input.size()
+
+    x = input.contiguous().view(b * c, w * h)
+
+    # x = x / 2
+
+    # Work around, torch.mm would generate some inf values.
+    # https://discuss.pytorch.org/t/gram-matrix-in-mixed-precision/166800/2
+    # x = torch.clamp(x, max=1.0e2, min=-1.0e2)
+    # x[x > 1.0e2] = 1.0e2
+    # x[x < -1.0e2] = -1.0e2
+
+    G = torch.mm(x, x.T)
+    G = torch.clamp(G, -64990.0, 64990.0)
+    # normalize by total elements
+    result = G.div(b * c * w * h)
+    return result
+
+
+
+def divisible(dim):
+    '''
+    Make width and height divisible by 32
+    '''
+    width, height = dim
+    return width - (width % 32), height - (height % 32)
+
+
+def resize_image(image, width=None, height=None, inter=cv2.INTER_AREA):
+    dim = None
+    h, w = image.shape[:2]
+
+    if width and height:
+        return cv2.resize(image, divisible((width, height)),  interpolation=inter)
+
+    if width is None and height is None:
+        return cv2.resize(image, divisible((w, h)),  interpolation=inter)
+
+    if width is None:
+        r = height / float(h)
+        dim = (int(w * r), height)
+
+    else:
+        r = width / float(w)
+        dim = (width, int(h * r))
+
+    return cv2.resize(image, divisible(dim), interpolation=inter)
+
+
+def normalize_input(images):
+    '''
+    [0, 255] -> [-1, 1]
+    '''
+    return images / 127.5 - 1.0
+
+
+def denormalize_input(images, dtype=None):
+    '''
+    [-1, 1] -> [0, 255]
+    '''
+    images = images * 127.5 + 127.5
+
+    if dtype is not None:
+        if isinstance(images, torch.Tensor):
+            images = images.type(dtype)
+        else:
+            # numpy.ndarray
+            images = images.astype(dtype)
+
+    return images
+
+
+def preprocess_images(images):
+    '''
+    Preprocess image for inference
+
+    @Arguments:
+        - images: np.ndarray
+
+    @Returns
+        - images: torch.tensor
+    '''
+    images = images.astype(np.float32)
+
+    # Normalize to [-1, 1]
+    images = normalize_input(images)
+    images = torch.from_numpy(images)
+
+    # Add batch dim
+    if len(images.shape) == 3:
+        images = images.unsqueeze(0)
+
+    # channel first
+    images = images.permute(0, 3, 1, 2)
+
+    return images
+
+def compute_data_mean(data_folder):
+    if not os.path.exists(data_folder):
+        raise FileNotFoundError(f'Folder {data_folder} does not exits')
+
+    image_files = os.listdir(data_folder)
+    total = np.zeros(3)
+
+    print(f"Compute mean (R, G, B) from {len(image_files)} images")
+
+    for img_file in tqdm(image_files):
+        path = os.path.join(data_folder, img_file)
+        image = cv2.imread(path)
+        total += image.mean(axis=(0, 1))
+
+    channel_mean = total / len(image_files)
+    mean = np.mean(channel_mean)
+
+    return mean - channel_mean[...,::-1]  # Convert to BGR for training
+
+
+if __name__ == '__main__':
+    t = torch.rand(2, 14, 32, 32)
+
+    with torch.autocast("cpu"):
+        print(gram(t))

utils/logger.py

@@ -0,0 +1,24 @@
+import logging
+
+
+def get_logger(path, *args, **kwargs):
+    # logger = logging.getLogger('train')
+    # logger.setLevel(logging.NOTSET)
+    # formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+    # # add filehandler
+    # fh = logging.FileHandler(path)
+    # fh.setLevel(logging.NOTSET)
+    # fh.setFormatter(formatter)
+    # ch = logging.StreamHandler()
+    # ch.setLevel(logging.ERROR)
+    # logger.addHandler(fh)
+    # logger.addHandler(ch)
+    # return logger
+    logging.basicConfig(format = '%(asctime)s %(message)s',
+                    datefmt = '%m/%d/%Y %I:%M:%S %p',
+                    handlers=[
+                        logging.FileHandler(path),
+                        logging.StreamHandler()
+                    ],
+                    level=logging.DEBUG)
+    return logging