| | import os |
| | import sys |
| |
|
| | from urllib.parse import urlparse |
| | import cv2 |
| | import numpy as np |
| | import torch |
| | from torch.hub import download_url_to_file, get_dir |
| |
|
| | LAMA_MODEL_URL = os.environ.get( |
| | "LAMA_MODEL_URL", |
| | "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt", |
| | ) |
| |
|
| |
|
| | def download_model(url=LAMA_MODEL_URL): |
| | parts = urlparse(url) |
| | hub_dir = get_dir() |
| | model_dir = os.path.join(hub_dir, "checkpoints") |
| | if not os.path.isdir(model_dir): |
| | os.makedirs(os.path.join(model_dir, "hub", "checkpoints")) |
| | filename = os.path.basename(parts.path) |
| | cached_file = os.path.join(model_dir, filename) |
| | if not os.path.exists(cached_file): |
| | sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file)) |
| | hash_prefix = None |
| | download_url_to_file(url, cached_file, hash_prefix, progress=True) |
| | return cached_file |
| |
|
| |
|
| | def ceil_modulo(x, mod): |
| | if x % mod == 0: |
| | return x |
| | return (x // mod + 1) * mod |
| |
|
| |
|
| | def numpy_to_bytes(image_numpy: np.ndarray) -> bytes: |
| | data = cv2.imencode(".jpg", image_numpy)[1] |
| | image_bytes = data.tobytes() |
| | return image_bytes |
| |
|
| |
|
| | def load_img(img_bytes, gray: bool = False): |
| | nparr = np.frombuffer(img_bytes, np.uint8) |
| | if gray: |
| | np_img = cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE) |
| | else: |
| | np_img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED) |
| | if len(np_img.shape) == 3 and np_img.shape[2] == 4: |
| | np_img = cv2.cvtColor(np_img, cv2.COLOR_BGRA2RGB) |
| | else: |
| | np_img = cv2.cvtColor(np_img, cv2.COLOR_BGR2RGB) |
| |
|
| | return np_img |
| |
|
| |
|
| | def norm_img(np_img): |
| | if len(np_img.shape) == 2: |
| | np_img = np_img[:, :, np.newaxis] |
| | np_img = np.transpose(np_img, (2, 0, 1)) |
| | np_img = np_img.astype("float32") / 255 |
| | return np_img |
| |
|
| |
|
| | def resize_max_size( |
| | np_img, size_limit: int, interpolation=cv2.INTER_CUBIC |
| | ) -> np.ndarray: |
| | |
| | h, w = np_img.shape[:2] |
| | if max(h, w) > size_limit: |
| | ratio = size_limit / max(h, w) |
| | new_w = int(w * ratio + 0.5) |
| | new_h = int(h * ratio + 0.5) |
| | return cv2.resize(np_img, dsize=(new_w, new_h), interpolation=interpolation) |
| | else: |
| | return np_img |
| |
|
| |
|
| | def pad_img_to_modulo(img, mod): |
| | channels, height, width = img.shape |
| | out_height = ceil_modulo(height, mod) |
| | out_width = ceil_modulo(width, mod) |
| | return np.pad( |
| | img, |
| | ((0, 0), (0, out_height - height), (0, out_width - width)), |
| | mode="symmetric", |
| | ) |
