218 lines
6.3 KiB
Python
218 lines
6.3 KiB
Python
import numpy as np
|
|
import cv2
|
|
import os
|
|
import math
|
|
import requests
|
|
from torch.hub import download_url_to_file, get_dir
|
|
from tqdm import tqdm
|
|
from urllib.parse import urlparse
|
|
import torch
|
|
import random
|
|
|
|
|
|
annotator_ckpts_path = os.path.join(os.path.dirname(__file__), "ckpts")
|
|
|
|
|
|
def HWC3(x):
|
|
assert x.dtype == np.uint8
|
|
if x.ndim == 2:
|
|
x = x[:, :, None]
|
|
assert x.ndim == 3
|
|
H, W, C = x.shape
|
|
assert C == 1 or C == 3 or C == 4
|
|
if C == 3:
|
|
return x
|
|
if C == 1:
|
|
return np.concatenate([x, x, x], axis=2)
|
|
if C == 4:
|
|
color = x[:, :, 0:3].astype(np.float32)
|
|
alpha = x[:, :, 3:4].astype(np.float32) / 255.0
|
|
y = color * alpha + 255.0 * (1.0 - alpha)
|
|
y = y.clip(0, 255).astype(np.uint8)
|
|
return y
|
|
|
|
|
|
def resize_image(input_image, resolution=None):
|
|
H, W, C = input_image.shape
|
|
H = float(H)
|
|
W = float(W)
|
|
k = 0
|
|
if resolution is not None:
|
|
k = float(resolution) / min(H, W)
|
|
H *= k
|
|
W *= k
|
|
H = int(np.round(H / 64.0)) * 64
|
|
W = int(np.round(W / 64.0)) * 64
|
|
img = cv2.resize(input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA)
|
|
return img
|
|
|
|
|
|
#####COPIED FROM BASICSR####
|
|
def sizeof_fmt(size, suffix='B'):
|
|
"""Get human readable file size.
|
|
|
|
Args:
|
|
size (int): File size.
|
|
suffix (str): Suffix. Default: 'B'.
|
|
|
|
Return:
|
|
str: Formatted file size.
|
|
"""
|
|
for unit in ['', 'K', 'M', 'G', 'T', 'P', 'E', 'Z']:
|
|
if abs(size) < 1024.0:
|
|
return f'{size:3.1f} {unit}{suffix}'
|
|
size /= 1024.0
|
|
return f'{size:3.1f} Y{suffix}'
|
|
|
|
|
|
def download_file_from_google_drive(file_id, save_path):
|
|
"""Download files from google drive.
|
|
|
|
Ref:
|
|
https://stackoverflow.com/questions/25010369/wget-curl-large-file-from-google-drive # noqa E501
|
|
|
|
Args:
|
|
file_id (str): File id.
|
|
save_path (str): Save path.
|
|
"""
|
|
|
|
session = requests.Session()
|
|
URL = 'https://docs.google.com/uc?export=download'
|
|
params = {'id': file_id}
|
|
|
|
response = session.get(URL, params=params, stream=True)
|
|
token = get_confirm_token(response)
|
|
if token:
|
|
params['confirm'] = token
|
|
response = session.get(URL, params=params, stream=True)
|
|
|
|
# get file size
|
|
response_file_size = session.get(URL, params=params, stream=True, headers={'Range': 'bytes=0-2'})
|
|
if 'Content-Range' in response_file_size.headers:
|
|
file_size = int(response_file_size.headers['Content-Range'].split('/')[1])
|
|
else:
|
|
file_size = None
|
|
|
|
save_response_content(response, save_path, file_size)
|
|
|
|
|
|
def get_confirm_token(response):
|
|
for key, value in response.cookies.items():
|
|
if key.startswith('download_warning'):
|
|
return value
|
|
return None
|
|
|
|
|
|
def save_response_content(response, destination, file_size=None, chunk_size=32768):
|
|
if file_size is not None:
|
|
pbar = tqdm(total=math.ceil(file_size / chunk_size), unit='chunk')
|
|
|
|
readable_file_size = sizeof_fmt(file_size)
|
|
else:
|
|
pbar = None
|
|
|
|
with open(destination, 'wb') as f:
|
|
downloaded_size = 0
|
|
for chunk in response.iter_content(chunk_size):
|
|
downloaded_size += chunk_size
|
|
if pbar is not None:
|
|
pbar.update(1)
|
|
pbar.set_description(f'Download {sizeof_fmt(downloaded_size)} / {readable_file_size}')
|
|
if chunk: # filter out keep-alive new chunks
|
|
f.write(chunk)
|
|
if pbar is not None:
|
|
pbar.close()
|
|
|
|
|
|
def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
|
|
"""Load file form http url, will download models if necessary.
|
|
|
|
Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
|
|
|
|
Args:
|
|
url (str): URL to be downloaded.
|
|
model_dir (str): The path to save the downloaded model. Should be a full path. If None, use pytorch hub_dir.
|
|
Default: None.
|
|
progress (bool): Whether to show the download progress. Default: True.
|
|
file_name (str): The downloaded file name. If None, use the file name in the url. Default: None.
|
|
|
|
Returns:
|
|
str: The path to the downloaded file.
|
|
"""
|
|
if model_dir is None: # use the pytorch hub_dir
|
|
hub_dir = get_dir()
|
|
model_dir = os.path.join(hub_dir, 'checkpoints')
|
|
|
|
os.makedirs(model_dir, exist_ok=True)
|
|
|
|
parts = urlparse(url)
|
|
filename = os.path.basename(parts.path)
|
|
if file_name is not None:
|
|
filename = file_name
|
|
cached_file = os.path.abspath(os.path.join(model_dir, filename))
|
|
if not os.path.exists(cached_file):
|
|
print(f'Downloading: "{url}" to {cached_file}\n')
|
|
download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
|
|
return cached_file
|
|
|
|
|
|
def load_state_dict(modelpath):
|
|
wrapper = torch.load(modelpath)
|
|
return wrapper["state_dict"] if "state_dict" in wrapper else wrapper
|
|
|
|
def nms(x, t, s):
|
|
x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
|
|
|
|
f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
|
|
f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
|
|
f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
|
|
f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
|
|
|
|
y = np.zeros_like(x)
|
|
|
|
for f in [f1, f2, f3, f4]:
|
|
np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
|
|
|
|
z = np.zeros_like(y, dtype=np.uint8)
|
|
z[y > t] = 255
|
|
return z
|
|
|
|
|
|
def make_noise_disk(H, W, C, F):
|
|
noise = np.random.uniform(low=0, high=1, size=((H // F) + 2, (W // F) + 2, C))
|
|
noise = cv2.resize(noise, (W + 2 * F, H + 2 * F), interpolation=cv2.INTER_CUBIC)
|
|
noise = noise[F: F + H, F: F + W]
|
|
noise -= np.min(noise)
|
|
noise /= np.max(noise)
|
|
if C == 1:
|
|
noise = noise[:, :, None]
|
|
return noise
|
|
|
|
|
|
def min_max_norm(x):
|
|
x -= np.min(x)
|
|
x /= np.maximum(np.max(x), 1e-5)
|
|
return x
|
|
|
|
|
|
def safe_step(x, step=2):
|
|
y = x.astype(np.float32) * float(step + 1)
|
|
y = y.astype(np.int32).astype(np.float32) / float(step)
|
|
return y
|
|
|
|
|
|
def img2mask(img, H, W, low=10, high=90):
|
|
assert img.ndim == 3 or img.ndim == 2
|
|
assert img.dtype == np.uint8
|
|
|
|
if img.ndim == 3:
|
|
y = img[:, :, random.randrange(0, img.shape[2])]
|
|
else:
|
|
y = img
|
|
|
|
y = cv2.resize(y, (W, H), interpolation=cv2.INTER_CUBIC)
|
|
|
|
if random.uniform(0, 1) < 0.5:
|
|
y = 255 - y
|
|
|
|
return y < np.percentile(y, random.randrange(low, high)) |