| import torch | |
| import torch.nn as nn | |
| import numpy as np | |
| import scipy.stats as stats | |
| import cv2 | |
| import json | |
| import patoolib | |
| import re | |
| from pathlib import Path | |
| from shutil import rmtree | |
| def weights_init(m): | |
| classname = m.__class__.__name__ | |
| if classname.find('Conv2d') != -1: | |
| nn.init.xavier_uniform_(m.weight.data) | |
| def weights_init_spectr(m): | |
| classname = m.__class__.__name__ | |
| if classname.find('Conv2d') != -1: | |
| nn.init.xavier_uniform_(m.weight_bar.data) | |
| def generate_mask(height, width, mu = 1, sigma = 0.0005, prob = 0.5, full = True, full_prob = 0.01): | |
| X = stats.truncnorm((0 - mu) / sigma, (1 - mu) / sigma, loc=mu, scale=sigma) | |
| if full: | |
| if (np.random.binomial(1, p = full_prob) == 1): | |
| return torch.ones(1, height, width).float() | |
| if np.random.binomial(1, p = prob) == 1: | |
| mask = torch.rand(1, height, width).ge(X.rvs(1)[0]).float() | |
| else: | |
| mask = torch.zeros(1, height, width).float() | |
| return mask | |
| def resize_pad(img, size = 512): | |
| if len(img.shape) == 2: | |
| img = np.expand_dims(img, 2) | |
| if img.shape[2] == 1: | |
| img = np.repeat(img, 3, 2) | |
| if img.shape[2] == 4: | |
| img = img[:, :, :3] | |
| pad = None | |
| if (img.shape[0] < img.shape[1]): | |
| height = img.shape[0] | |
| ratio = height / size | |
| width = int(np.ceil(img.shape[1] / ratio)) | |
| img = cv2.resize(img, (width, size), interpolation = cv2.INTER_AREA) | |
| new_width = width | |
| while (new_width % 32 != 0): | |
| new_width += 1 | |
| pad = (0, new_width - width) | |
| img = np.pad(img, ((0, 0), (0, pad[1]), (0, 0)), 'maximum') | |
| else: | |
| width = img.shape[1] | |
| ratio = width / size | |
| height = int(np.ceil(img.shape[0] / ratio)) | |
| img = cv2.resize(img, (size, height), interpolation = cv2.INTER_AREA) | |
| new_height = height | |
| while (new_height % 32 != 0): | |
| new_height += 1 | |
| pad = (new_height - height, 0) | |
| img = np.pad(img, ((0, pad[0]), (0, 0), (0, 0)), 'maximum') | |
| if (img.dtype == 'float32'): | |
| np.clip(img, 0, 1, out = img) | |
| return img, pad | |
| def open_json(file): | |
| with open(file) as json_file: | |
| data = json.load(json_file) | |
| return data | |
| def extract_cbr(file, out_dir): | |
| patoolib.extract_archive(file, outdir = out_dir, verbosity = 1) | |
| def create_cbz(file_path, files): | |
| patoolib.create_archive(file_path, files, verbosity = 1) | |
| def subfolder_image_search(start_folder): | |
| return [x.as_posix() for x in Path(".").rglob("*.[pPjJ][nNpP][gG]")] | |
| def remove_folder(folder_path): | |
| rmtree(folder_path) | |
| def sorted_alphanumeric(data): | |
| convert = lambda text: int(text) if text.isdigit() else text.lower() | |
| alphanum_key = lambda key: [ convert(c) for c in re.split('([0-9]+)', key) ] | |
| return sorted(data, key=alphanum_key) |