utils.py

import torch
import torch.nn as nn
from enum import Enum

import base64
import json
from io import BytesIO
from PIL import Image
import requests
import re
from copy import deepcopy

class ImageType(Enum):
    REAL_UP_L = 0
    REAL_UP_R = 1
    REAL_DOWN_R = 2
    REAL_DOWN_L = 3
    FAKE = 4


def crop_image_part(image: torch.Tensor,
                    part: ImageType) -> torch.Tensor:
    size = image.shape[2] // 2

    if part == ImageType.REAL_UP_L:
        return image[:, :, :size, :size]

    elif part == ImageType.REAL_UP_R:
        return image[:, :, :size, size:]

    elif part == ImageType.REAL_DOWN_L:
        return image[:, :, size:, :size]

    elif part == ImageType.REAL_DOWN_R:
        return image[:, :, size:, size:]

    else:
        raise ValueError('invalid part')


def init_weights(module: nn.Module):
    if isinstance(module, nn.Conv2d):
        torch.nn.init.normal_(module.weight, 0.0, 0.02)

    if isinstance(module, nn.BatchNorm2d):
        torch.nn.init.normal_(module.weight, 1.0, 0.02)
        module.bias.data.fill_(0)

def load_image_from_local(image_path, image_resize=None):
    image = Image.open(image_path)

    if isinstance(image_resize, tuple):
        image = image.resize(image_resize)
    return image

def load_image_from_url(image_url, rgba_mode=False, image_resize=None, default_image=None):
    try:
        image = Image.open(requests.get(image_url, stream=True).raw)

        if rgba_mode:
            image = image.convert("RGBA")

        if isinstance(image_resize, tuple):
            image = image.resize(image_resize)

    except Exception as e:
        image = None
        if default_image:
            image = load_image_from_local(default_image, image_resize=image_resize)

    return image

def image_to_base64(image_array):
    buffered = BytesIO()
    image_array.save(buffered, format="PNG")
    image_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
    return f"data:image/png;base64, {image_b64}"


def copy_G_params(model):
    flatten = deepcopy(list(p.data for p in model.parameters()))
    return flatten


def load_params(model, new_param):
    for p, new_p in zip(model.parameters(), new_param):
        p.data.copy_(new_p)